{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff734a78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | default_exp _application.app"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5096aab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "import asyncio\n",
    "import functools\n",
    "import inspect\n",
    "import json\n",
    "import types\n",
    "from asyncio import iscoroutinefunction  # do not use the version from inspect\n",
    "from collections import namedtuple\n",
    "from contextlib import asynccontextmanager\n",
    "from copy import deepcopy\n",
    "from datetime import datetime, timedelta\n",
    "from functools import wraps\n",
    "from inspect import signature\n",
    "from pathlib import Path\n",
    "from typing import *\n",
    "from unittest.mock import AsyncMock, MagicMock\n",
    "\n",
    "import anyio\n",
    "from pydantic import BaseModel\n",
    "\n",
    "from fastkafka._components.aiokafka_consumer_loop import (\n",
    "    aiokafka_consumer_loop,\n",
    "    sanitize_kafka_config,\n",
    ")\n",
    "from fastkafka._components.asyncapi import (\n",
    "    ConsumeCallable,\n",
    "    ContactInfo,\n",
    "    KafkaBroker,\n",
    "    KafkaBrokers,\n",
    "    KafkaServiceInfo,\n",
    "    export_async_spec,\n",
    ")\n",
    "\n",
    "import fastkafka._aiokafka_imports\n",
    "from fastkafka._aiokafka_imports import AIOKafkaConsumer, AIOKafkaProducer\n",
    "from fastkafka._components.benchmarking import _benchmark\n",
    "from fastkafka._components.logger import get_logger\n",
    "from fastkafka._components.meta import delegates, export, filter_using_signature, patch\n",
    "from fastkafka._components.producer_decorator import ProduceCallable, producer_decorator\n",
    "from fastkafka._components.task_streaming import StreamExecutor\n",
    "from fastkafka._components.helpers import remove_suffix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85f14f6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "if TYPE_CHECKING:\n",
    "    from fastapi import FastAPI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdedeee5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import __main__\n",
    "\n",
    "import os\n",
    "import shutil\n",
    "import unittest.mock\n",
    "from contextlib import asynccontextmanager\n",
    "\n",
    "import asyncer\n",
    "import pytest\n",
    "import requests\n",
    "import uvicorn\n",
    "import yaml\n",
    "from fastapi import FastAPI\n",
    "from pydantic import EmailStr, Field, HttpUrl\n",
    "\n",
    "from fastkafka._components.helpers import true_after\n",
    "from fastkafka._components.logger import suppress_timestamps\n",
    "from fastkafka._server import run_in_process\n",
    "from fastkafka.encoder import avro_decoder, avro_encoder, json_decoder, json_encoder\n",
    "from fastkafka.testing import ApacheKafkaBroker, Tester, mock_AIOKafkaProducer_send"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b33a28e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "logger = get_logger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cf16b2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: ok\n"
     ]
    }
   ],
   "source": [
    "suppress_timestamps()\n",
    "logger = get_logger(__name__, level=20)\n",
    "logger.info(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9177ac1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | notest\n",
    "\n",
    "# allows async calls in notebooks\n",
    "\n",
    "import nest_asyncio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e0f175c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | notest\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "211534a4",
   "metadata": {},
   "source": [
    "### Constructor utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f978e721",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@delegates(fastkafka._aiokafka_imports.AIOKafkaConsumer, but=[\"bootstrap_servers\"])\n",
    "@delegates(fastkafka._aiokafka_imports.AIOKafkaProducer, but=[\"bootstrap_servers\"], keep=True)\n",
    "def _get_kafka_config(\n",
    "    bootstrap_servers_id: str = \"localhost\",\n",
    "    **kwargs: Any,\n",
    ") -> Dict[str, Any]:\n",
    "    \"\"\"Get kafka config\"\"\"\n",
    "    allowed_keys = set(signature(_get_kafka_config).parameters.keys())\n",
    "    if not set(kwargs.keys()) <= allowed_keys:\n",
    "        unallowed_keys = \", \".join(\n",
    "            sorted([f\"'{x}'\" for x in set(kwargs.keys()).difference(allowed_keys)])\n",
    "        )\n",
    "        raise ValueError(f\"Unallowed key arguments passed: {unallowed_keys}\")\n",
    "    retval = kwargs.copy()\n",
    "\n",
    "    # todo: check this values\n",
    "    config_defaults = {\n",
    "        \"bootstrap_servers_id\": bootstrap_servers_id,\n",
    "        \"auto_offset_reset\": \"earliest\",\n",
    "        \"max_poll_records\": 100,\n",
    "    }\n",
    "    for key, value in config_defaults.items():\n",
    "        if key not in retval:\n",
    "            retval[key] = value\n",
    "\n",
    "    return retval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cdf9e16",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert _get_kafka_config() == {\n",
    "    \"bootstrap_servers_id\": \"localhost\",\n",
    "    \"auto_offset_reset\": \"earliest\",\n",
    "    \"max_poll_records\": 100,\n",
    "}\n",
    "\n",
    "assert _get_kafka_config(max_poll_records=1_000) == {\n",
    "    \"bootstrap_servers_id\": \"localhost\",\n",
    "    \"auto_offset_reset\": \"earliest\",\n",
    "    \"max_poll_records\": 1_000,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34b17b9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "with pytest.raises(ValueError) as e:\n",
    "    _get_kafka_config(random_key=1_000, whatever=\"whocares\")\n",
    "assert e.value.args == (\"Unallowed key arguments passed: 'random_key', 'whatever'\",)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3477cb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _get_kafka_brokers(\n",
    "    kafka_brokers: Optional[Union[Dict[str, Any], KafkaBrokers]] = None\n",
    ") -> KafkaBrokers:\n",
    "    \"\"\"Get Kafka brokers\n",
    "\n",
    "    Args:\n",
    "        kafka_brokers: Kafka brokers\n",
    "\n",
    "    \"\"\"\n",
    "    if kafka_brokers is None:\n",
    "        retval: KafkaBrokers = KafkaBrokers(\n",
    "            brokers={\n",
    "                \"localhost\": KafkaBroker(  # type: ignore\n",
    "                    url=\"https://localhost\",\n",
    "                    description=\"Local (dev) Kafka broker\",\n",
    "                    port=\"9092\",\n",
    "                    grouping=\"localhost\",\n",
    "                )\n",
    "            }\n",
    "        )\n",
    "    else:\n",
    "        if isinstance(kafka_brokers, KafkaBrokers):\n",
    "            return kafka_brokers\n",
    "\n",
    "        retval = KafkaBrokers(\n",
    "            brokers={\n",
    "                k: (\n",
    "                    [\n",
    "                        KafkaBroker.model_validate_json(\n",
    "                            unwrapped_v.model_dump_json()\n",
    "                            if hasattr(unwrapped_v, \"model_dump_json\")\n",
    "                            else json.dumps(unwrapped_v)\n",
    "                        )\n",
    "                        for unwrapped_v in v\n",
    "                    ]\n",
    "                    if isinstance(v, list)\n",
    "                    else KafkaBroker.model_validate_json(\n",
    "                        v.model_dump_json() if hasattr(v, \"model_dump_json\") else json.dumps(v)\n",
    "                    )\n",
    "                )\n",
    "                for k, v in kafka_brokers.items()\n",
    "            }\n",
    "        )\n",
    "\n",
    "    return retval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4449bac2",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert (\n",
    "    _get_kafka_brokers(None).model_dump_json()\n",
    "    == '{\"brokers\": {\"localhost\": {\"url\": \"https://localhost\", \"description\": \"Local (dev) Kafka broker\", \"protocol\": \"kafka\", \"variables\": {\"port\": {\"default\": \"9092\"}}}}}'\n",
    ")\n",
    "\n",
    "assert (\n",
    "    _get_kafka_brokers(dict(localhost=dict(url=\"localhost\"))).model_dump_json()\n",
    "    == '{\"brokers\": {\"localhost\": {\"url\": \"localhost\", \"description\": \"Kafka broker\", \"protocol\": \"kafka\", \"variables\": {\"port\": {\"default\": \"9092\"}}}}}'\n",
    ")\n",
    "\n",
    "assert (\n",
    "    _get_kafka_brokers(\n",
    "        dict(localhost=dict(url=\"localhost\"), staging=dict(url=\"staging.airt.ai\"))\n",
    "    ).model_dump_json()\n",
    "    == '{\"brokers\": {\"localhost\": {\"url\": \"localhost\", \"description\": \"Kafka broker\", \"protocol\": \"kafka\", \"variables\": {\"port\": {\"default\": \"9092\"}}}, \"staging\": {\"url\": \"staging.airt.ai\", \"description\": \"Kafka broker\", \"protocol\": \"kafka\", \"variables\": {\"port\": {\"default\": \"9092\"}}}}}'\n",
    ")\n",
    "\n",
    "assert (\n",
    "    _get_kafka_brokers(\n",
    "        dict(\n",
    "            localhost=[dict(url=\"localhost123\"), dict(url=\"localhost321\")],\n",
    "            staging=dict(url=\"staging.airt.ai\"),\n",
    "        )\n",
    "    ).model_dump_json()\n",
    "    == '{\"brokers\": {\"localhost-bootstrap-server-0\": {\"url\": \"localhost123\", \"description\": \"Kafka broker\", \"protocol\": \"kafka\", \"variables\": {\"port\": {\"default\": \"9092\"}}}, \"localhost-bootstrap-server-1\": {\"url\": \"localhost321\", \"description\": \"Kafka broker\", \"protocol\": \"kafka\", \"variables\": {\"port\": {\"default\": \"9092\"}}}, \"staging\": {\"url\": \"staging.airt.ai\", \"description\": \"Kafka broker\", \"protocol\": \"kafka\", \"variables\": {\"port\": {\"default\": \"9092\"}}}}}'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b175647e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "def _get_broker_addr_list(\n",
    "    brokers: Union[List[KafkaBroker], KafkaBroker]\n",
    ") -> Union[str, List[str]]:\n",
    "    if isinstance(brokers, list):\n",
    "        return [f\"{broker.url}:{broker.port}\" for broker in brokers]\n",
    "    else:\n",
    "        return f\"{brokers.url}:{brokers.port}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cf5447e",
   "metadata": {},
   "outputs": [],
   "source": [
    "brokers_config = _get_kafka_brokers(\n",
    "    dict(\n",
    "        localhost=[dict(url=\"localhost123\"), dict(url=\"localhost321\")],\n",
    "        staging=dict(url=\"staging.airt.ai\"),\n",
    "    )\n",
    ")\n",
    "\n",
    "assert _get_broker_addr_list(brokers_config.brokers[\"localhost\"]) == ['localhost123:9092', 'localhost321:9092']\n",
    "assert _get_broker_addr_list(brokers_config.brokers[\"staging\"]) == 'staging.airt.ai:9092'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f6c1f40",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _get_topic_name(\n",
    "    topic_callable: Union[ConsumeCallable, ProduceCallable], prefix: str = \"on_\"\n",
    ") -> str:\n",
    "    \"\"\"Get topic name\n",
    "    Args:\n",
    "        topic_callable: a function\n",
    "        prefix: prefix of the name of the function followed by the topic name\n",
    "\n",
    "    Returns:\n",
    "        The name of the topic\n",
    "    \"\"\"\n",
    "    topic = topic_callable.__name__\n",
    "    if not topic.startswith(prefix) or len(topic) <= len(prefix):\n",
    "        raise ValueError(f\"Function name '{topic}' must start with {prefix}\")\n",
    "    topic = topic[len(prefix) :]\n",
    "\n",
    "    return topic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec9b3689",
   "metadata": {},
   "outputs": [],
   "source": [
    "def on_topic_name_1():\n",
    "    pass\n",
    "\n",
    "\n",
    "assert _get_topic_name(on_topic_name_1) == \"topic_name_1\"\n",
    "\n",
    "assert _get_topic_name(on_topic_name_1, prefix=\"on_topic_\") == \"name_1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab9e9f96",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _get_contact_info(\n",
    "    name: str = \"Author\",\n",
    "    url: str = \"https://www.google.com\",\n",
    "    email: str = \"noreply@gmail.com\",\n",
    ") -> ContactInfo:\n",
    "    return ContactInfo(name=name, url=url, email=email)  # type: ignore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7e311f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert _get_contact_info() == ContactInfo(\n",
    "    name=\"Author\",\n",
    "    url=HttpUrl(url=\"https://www.google.com\"),\n",
    "    email=\"noreply@gmail.com\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57eeeb84",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | exporti\n",
    "\n",
    "I = TypeVar(\"I\", bound=BaseModel)\n",
    "O = TypeVar(\"O\", BaseModel, Awaitable[BaseModel])\n",
    "\n",
    "F = TypeVar(\"F\", bound=Callable)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49c37353",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@export(\"fastkafka\")\n",
    "class FastKafka:\n",
    "    @delegates(_get_kafka_config)\n",
    "    def __init__(\n",
    "        self,\n",
    "        *,\n",
    "        title: Optional[str] = None,\n",
    "        description: Optional[str] = None,\n",
    "        version: Optional[str] = None,\n",
    "        contact: Optional[Dict[str, str]] = None,\n",
    "        kafka_brokers: Optional[Dict[str, Any]] = None,\n",
    "        root_path: Optional[Union[Path, str]] = None,\n",
    "        lifespan: Optional[Callable[[\"FastKafka\"], AsyncContextManager[None]]] = None,\n",
    "        **kwargs: Any,\n",
    "    ):\n",
    "        \"\"\"Creates FastKafka application\n",
    "\n",
    "        Args:\n",
    "            title: optional title for the documentation. If None,\n",
    "                the title will be set to empty string\n",
    "            description: optional description for the documentation. If\n",
    "                None, the description will be set to empty string\n",
    "            version: optional version for the documentation. If None,\n",
    "                the version will be set to empty string\n",
    "            contact: optional contact for the documentation. If None, the\n",
    "                contact will be set to placeholder values:\n",
    "                name='Author' url=HttpUrl('https://www.google.com', ) email='noreply@gmail.com'\n",
    "            kafka_brokers: dictionary describing kafka brokers used for setting\n",
    "                the bootstrap server when running the applicationa and for\n",
    "                generating documentation. Defaults to\n",
    "                    {\n",
    "                        \"localhost\": {\n",
    "                            \"url\": \"localhost\",\n",
    "                            \"description\": \"local kafka broker\",\n",
    "                            \"port\": \"9092\",\n",
    "                        }\n",
    "                    }\n",
    "            root_path: path to where documentation will be created\n",
    "            lifespan: asynccontextmanager that is used for setting lifespan hooks.\n",
    "                __aenter__ is called before app start and __aexit__ after app stop.\n",
    "                The lifespan is called whe application is started as async context\n",
    "                manager, e.g.:`async with kafka_app...`\n",
    "\n",
    "        \"\"\"\n",
    "\n",
    "        # this is needed for documentation generation\n",
    "        self._title = title if title is not None else \"\"\n",
    "        self._description = description if description is not None else \"\"\n",
    "        self._version = version if version is not None else \"\"\n",
    "        if contact is not None:\n",
    "            self._contact_info = _get_contact_info(**contact)\n",
    "        else:\n",
    "            self._contact_info = _get_contact_info()\n",
    "\n",
    "        self._kafka_service_info = KafkaServiceInfo(\n",
    "            title=self._title,\n",
    "            version=self._version,\n",
    "            description=self._description,\n",
    "            contact=self._contact_info,\n",
    "        )\n",
    "\n",
    "        if kafka_brokers is None:\n",
    "            kafka_brokers = {\n",
    "                \"localhost\": {\n",
    "                    \"url\": \"localhost\",\n",
    "                    \"description\": \"local kafka broker\",\n",
    "                    \"port\": \"9092\",\n",
    "                }\n",
    "            }\n",
    "\n",
    "        self._kafka_brokers = _get_kafka_brokers(kafka_brokers)\n",
    "\n",
    "        self._override_brokers: List[KafkaBrokers] = []\n",
    "\n",
    "        self._root_path = Path(\".\") if root_path is None else Path(root_path)\n",
    "        self._root_path.mkdir(exist_ok=True, parents=True)\n",
    "\n",
    "        self._asyncapi_path = self._root_path / \"asyncapi\"\n",
    "\n",
    "        # this is used as default parameters for creating AIOProducer and AIOConsumer objects\n",
    "        self._kafka_config = _get_kafka_config(**kwargs)\n",
    "\n",
    "        #\n",
    "        self._consumers_store: Dict[\n",
    "            str,\n",
    "            Tuple[\n",
    "                ConsumeCallable,\n",
    "                Callable[[bytes, Type[BaseModel]], Any],\n",
    "                Union[str, StreamExecutor, None],\n",
    "                Optional[KafkaBrokers],\n",
    "                Dict[str, Any],\n",
    "            ],\n",
    "        ] = {}\n",
    "\n",
    "        self._producers_store: Dict[  # type: ignore\n",
    "            str,\n",
    "            Tuple[\n",
    "                ProduceCallable,\n",
    "                fastkafka._aiokafka_imports.AIOKafkaProducer,\n",
    "                Optional[KafkaBrokers],\n",
    "                Dict[str, Any],\n",
    "            ],\n",
    "        ] = {}\n",
    "\n",
    "        self._producers_list: List[fastkafka._aiokafka_imports.AIOKafkaProducer] = []  # type: ignore\n",
    "\n",
    "        self.benchmark_results: Dict[str, Dict[str, Any]] = {}\n",
    "\n",
    "        # background tasks\n",
    "        self._scheduled_bg_tasks: List[Callable[..., Coroutine[Any, Any, Any]]] = []\n",
    "        self._bg_task_group_generator: Optional[anyio.abc.TaskGroup] = None\n",
    "        self._bg_tasks_group: Optional[anyio.abc.TaskGroup] = None\n",
    "\n",
    "        # todo: use this for errrors\n",
    "        self._on_error_topic: Optional[str] = None\n",
    "\n",
    "        self.lifespan = lifespan\n",
    "        self.lifespan_ctx: Optional[AsyncContextManager[None]] = None\n",
    "\n",
    "        self._is_started: bool = False\n",
    "        self._is_shutting_down: bool = False\n",
    "        self._kafka_consumer_tasks: List[asyncio.Task[Any]] = []\n",
    "        self._kafka_producer_tasks: List[asyncio.Task[Any]] = []\n",
    "        self._running_bg_tasks: List[asyncio.Task[Any]] = []\n",
    "        self.run = False\n",
    "\n",
    "        # testing functions\n",
    "        self.AppMocks = None\n",
    "        self.mocks = None\n",
    "        self.awaited_mocks = None\n",
    "\n",
    "    @property\n",
    "    def is_started(self) -> bool:\n",
    "        \"\"\"Property indicating whether the FastKafka object is started.\n",
    "\n",
    "        The is_started property indicates if the FastKafka object is currently \n",
    "        in a started state. This implies that all background tasks, producers, \n",
    "        and consumers have been initiated, and the object is successfully connected \n",
    "        to the Kafka broker. \n",
    "\n",
    "        Returns:\n",
    "            bool: True if the object is started, False otherwise.\n",
    "        \"\"\"\n",
    "        return self._is_started\n",
    "\n",
    "    def set_kafka_broker(self, kafka_broker_name: str) -> None:\n",
    "        \"\"\"\n",
    "        Sets the Kafka broker to start FastKafka with\n",
    "\n",
    "        Args:\n",
    "            kafka_broker_name: The name of the Kafka broker to start FastKafka\n",
    "\n",
    "        Raises:\n",
    "            ValueError: If the provided kafka_broker_name is not found in dictionary of kafka_brokers\n",
    "        \"\"\"\n",
    "\n",
    "        if kafka_broker_name not in self._kafka_brokers.brokers:\n",
    "            raise ValueError(\n",
    "                f\"Given kafka_broker_name '{kafka_broker_name}' is not found in kafka_brokers, available options are {self._kafka_brokers.brokers.keys()}\"\n",
    "            )\n",
    "\n",
    "        self._kafka_config[\"bootstrap_servers_id\"] = kafka_broker_name\n",
    "\n",
    "    async def __aenter__(self) -> \"FastKafka\":\n",
    "        if self.lifespan is not None:\n",
    "            self.lifespan_ctx = self.lifespan(self)\n",
    "            await self.lifespan_ctx.__aenter__()\n",
    "        await self._start()\n",
    "        return self\n",
    "\n",
    "    async def __aexit__(\n",
    "        self,\n",
    "        exc_type: Optional[Type[BaseException]],\n",
    "        exc: Optional[BaseException],\n",
    "        tb: Optional[types.TracebackType],\n",
    "    ) -> None:\n",
    "        await self._stop()\n",
    "        if self.lifespan_ctx is not None:\n",
    "            await self.lifespan_ctx.__aexit__(exc_type, exc, tb)\n",
    "\n",
    "    async def _start(self) -> None:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    async def _stop(self) -> None:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def consumes(\n",
    "        self,\n",
    "        topic: Optional[str] = None,\n",
    "        decoder: str = \"json\",\n",
    "        *,\n",
    "        prefix: str = \"on_\",\n",
    "        brokers: Optional[KafkaBrokers] = None,\n",
    "        description: Optional[str] = None,\n",
    "        **kwargs: Dict[str, Any],\n",
    "    ) -> ConsumeCallable:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def produces(\n",
    "        self,\n",
    "        topic: Optional[str] = None,\n",
    "        encoder: str = \"json\",\n",
    "        *,\n",
    "        prefix: str = \"to_\",\n",
    "        brokers: Optional[KafkaBrokers] = None,\n",
    "        description: Optional[str] = None,\n",
    "        **kwargs: Dict[str, Any],\n",
    "    ) -> ProduceCallable:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def benchmark(\n",
    "        self,\n",
    "        interval: Union[int, timedelta] = 1,\n",
    "        *,\n",
    "        sliding_window_size: Optional[int] = None,\n",
    "    ) -> Callable[[F], F]:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def run_in_background(\n",
    "        self,\n",
    "    ) -> Callable[[], Any]:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def _populate_consumers(\n",
    "        self,\n",
    "        is_shutting_down_f: Callable[[], bool],\n",
    "    ) -> None:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def get_topics(self) -> Iterable[str]:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    async def _populate_producers(self) -> None:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    async def _populate_bg_tasks(self) -> None:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def create_docs(self) -> None:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def create_mocks(self) -> None:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    async def _shutdown_consumers(self) -> None:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    async def _shutdown_producers(self) -> None:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    async def _shutdown_bg_tasks(self) -> None:\n",
    "        raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c25f82b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert FastKafka.__module__ == \"fastkafka\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "894af799",
   "metadata": {},
   "outputs": [],
   "source": [
    "kafka_app = FastKafka(\n",
    "    kafka_brokers=dict(\n",
    "        localhost=[dict(url=\"localhost\", port=\"9092\"), dict(url=\"localhost\", port=\"9093\")]\n",
    "    )\n",
    ")\n",
    "assert kafka_app._kafka_brokers == KafkaBrokers(\n",
    "    brokers={\n",
    "        \"localhost\": [\n",
    "            KafkaBroker(\n",
    "                url=\"localhost\",\n",
    "                description=\"Kafka broker\",\n",
    "                port=\"9092\",\n",
    "                protocol=\"kafka\",\n",
    "                security=None,\n",
    "            ),\n",
    "            KafkaBroker(\n",
    "                url=\"localhost\",\n",
    "                description=\"Kafka broker\",\n",
    "                port=\"9093\",\n",
    "                protocol=\"kafka\",\n",
    "                security=None,\n",
    "            ),\n",
    "        ]\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "997906d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "kafka_app = FastKafka()\n",
    "assert kafka_app._kafka_brokers == KafkaBrokers(\n",
    "    brokers={\n",
    "        \"localhost\": KafkaBroker(\n",
    "            url=\"localhost\",\n",
    "            description=\"local kafka broker\",\n",
    "            port=\"9092\",\n",
    "            protocol=\"kafka\",\n",
    "            security=None,\n",
    "        )\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cfbe17d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_testing_app(\n",
    "    *, root_path: str = \"/tmp/000_FastKafka\", bootstrap_servers: Optional[str] = None\n",
    "):\n",
    "    if Path(root_path).exists():\n",
    "        shutil.rmtree(root_path)\n",
    "\n",
    "    host, port = None, None\n",
    "    if bootstrap_servers is not None:\n",
    "        host, port = bootstrap_servers.split(\":\")\n",
    "\n",
    "    kafka_app = FastKafka(\n",
    "        kafka_brokers={\n",
    "            \"localhost\": {\n",
    "                \"url\": host if host is not None else \"localhost\",\n",
    "                \"name\": \"development\",\n",
    "                \"description\": \"Local (dev) Kafka broker\",\n",
    "                \"port\": port if port is not None else \"9092\",\n",
    "            }\n",
    "        },\n",
    "        root_path=root_path,\n",
    "    )\n",
    "    kafka_app.set_kafka_broker(kafka_broker_name=\"localhost\")\n",
    "\n",
    "    return kafka_app"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66237424",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<fastkafka.FastKafka>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "app = create_testing_app()\n",
    "assert Path(\"/tmp/000_FastKafka\").exists()\n",
    "app"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bebbf80",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _get_decoder_fn(decoder: str) -> Callable[[bytes, Type[BaseModel]], Any]:\n",
    "    \"\"\"\n",
    "    Imports and returns decoder function based on input\n",
    "    \"\"\"\n",
    "    if decoder == \"json\":\n",
    "        from fastkafka._components.encoder.json import json_decoder\n",
    "\n",
    "        return json_decoder\n",
    "    elif decoder == \"avro\":\n",
    "        try:\n",
    "            from fastkafka._components.encoder.avro import avro_decoder\n",
    "        except ModuleNotFoundError:\n",
    "            raise ModuleNotFoundError(\n",
    "                \"Unable to import avro packages. Please install FastKafka using the command 'fastkafka[avro]'\"\n",
    "            )\n",
    "        return avro_decoder\n",
    "    else:\n",
    "        raise ValueError(f\"Unknown decoder - {decoder}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6715c512",
   "metadata": {},
   "outputs": [],
   "source": [
    "actual = _get_decoder_fn(\"json\")\n",
    "assert actual == json_decoder\n",
    "\n",
    "actual = _get_decoder_fn(\"avro\")\n",
    "assert actual == avro_decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a38da3a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _prepare_and_check_brokers(\n",
    "    app: FastKafka, kafka_brokers: Optional[Union[Dict[str, Any], KafkaBrokers]]\n",
    ") -> Optional[KafkaBrokers]:\n",
    "    if kafka_brokers is not None:\n",
    "        prepared_brokers = _get_kafka_brokers(kafka_brokers)\n",
    "        if prepared_brokers.brokers.keys() != app._kafka_brokers.brokers.keys():\n",
    "            raise ValueError(\n",
    "                f\"To override application default brokers, you must define all of the broker options. Default defined: {set(app._kafka_brokers.brokers.keys())}, override defined: {set(prepared_brokers.brokers.keys())}\"\n",
    "            )\n",
    "        return prepared_brokers\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4b8b1ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _resolve_key(key: str, dictionary: Dict[str, Any]) -> str:\n",
    "    i = 0\n",
    "    resolved_key = f\"{key}_{i}\"\n",
    "    while resolved_key in dictionary:\n",
    "        i += 1\n",
    "        resolved_key = f\"{key}_{i}\"\n",
    "    return resolved_key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cddd5c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "@delegates(fastkafka._aiokafka_imports.AIOKafkaConsumer)\n",
    "def consumes(\n",
    "    self: FastKafka,\n",
    "    topic: Optional[str] = None,\n",
    "    decoder: Union[str, Callable[[bytes, Type[BaseModel]], Any]] = \"json\",\n",
    "    *,\n",
    "    executor: Union[str, StreamExecutor, None] = None,\n",
    "    brokers: Optional[Union[Dict[str, Any], KafkaBrokers]] = None,\n",
    "    prefix: str = \"on_\",\n",
    "    description: Optional[str] = None,\n",
    "    **kwargs: Dict[str, Any],\n",
    ") -> Callable[[ConsumeCallable], ConsumeCallable]:\n",
    "    \"\"\"Decorator registering the callback called when a message is received in a topic.\n",
    "\n",
    "    This function decorator is also responsible for registering topics for AsyncAPI specificiation and documentation.\n",
    "\n",
    "    Args:\n",
    "        topic: Kafka topic that the consumer will subscribe to and execute the\n",
    "            decorated function when it receives a message from the topic,\n",
    "            default: None. If the topic is not specified, topic name will be\n",
    "            inferred from the decorated function name by stripping the defined prefix\n",
    "        decoder: Decoder to use to decode messages consumed from the topic,\n",
    "                default: json - By default, it uses json decoder to decode\n",
    "                bytes to json string and then it creates instance of pydantic\n",
    "                BaseModel. It also accepts custom decoder function.\n",
    "        executor: Type of executor to choose for consuming tasks. Avaliable options\n",
    "                are \"SequentialExecutor\" and \"DynamicTaskExecutor\". The default option is\n",
    "                \"SequentialExecutor\" which will execute the consuming tasks sequentially.\n",
    "                If the consuming tasks have high latency it is recommended to use\n",
    "                \"DynamicTaskExecutor\" which will wrap the consuming functions into tasks\n",
    "                and run them in on asyncio loop in background. This comes with a cost of\n",
    "                increased overhead so use it only in cases when your consume functions have\n",
    "                high latency such as database queries or some other type of networking.\n",
    "        prefix: Prefix stripped from the decorated function to define a topic name\n",
    "                if the topic argument is not passed, default: \"on_\". If the decorated\n",
    "                function name is not prefixed with the defined prefix and topic argument\n",
    "                is not passed, then this method will throw ValueError\n",
    "        brokers: Optional argument specifying multiple broker clusters for consuming\n",
    "                messages from different Kafka clusters in FastKafka.\n",
    "        description: Optional description of the consuming function async docs.\n",
    "                If not provided, consuming function __doc__ attr will be used.\n",
    "\n",
    "    Returns:\n",
    "        A function returning the same function\n",
    "\n",
    "    Throws:\n",
    "        ValueError\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def _decorator(\n",
    "        on_topic: ConsumeCallable,\n",
    "        topic: Optional[str] = topic,\n",
    "        decoder: Union[str, Callable[[bytes, Type[BaseModel]], Any]] = decoder,\n",
    "        executor: Union[str, StreamExecutor, None] = executor,\n",
    "        brokers: Optional[Union[Dict[str, Any], KafkaBrokers]] = brokers,\n",
    "        description: Optional[str] = description,\n",
    "        kwargs: Dict[str, Any] = kwargs,\n",
    "    ) -> ConsumeCallable:\n",
    "        topic_resolved: str = (\n",
    "            _get_topic_name(topic_callable=on_topic, prefix=prefix)\n",
    "            if topic is None\n",
    "            else topic\n",
    "        )\n",
    "\n",
    "        decoder_fn = _get_decoder_fn(decoder) if isinstance(decoder, str) else decoder\n",
    "\n",
    "        prepared_broker = _prepare_and_check_brokers(self, brokers)\n",
    "        if prepared_broker is not None:\n",
    "            self._override_brokers.append(prepared_broker.brokers)  # type: ignore\n",
    "        else:\n",
    "            prepared_broker = self._kafka_brokers\n",
    "\n",
    "        if description is not None:\n",
    "            setattr(on_topic, \"description\", description)\n",
    "\n",
    "        self._consumers_store[_resolve_key(topic_resolved, self._consumers_store)] = (\n",
    "            on_topic,\n",
    "            decoder_fn,\n",
    "            executor,\n",
    "            prepared_broker,\n",
    "            kwargs,\n",
    "        )\n",
    "        setattr(self, on_topic.__name__, on_topic)\n",
    "        return on_topic\n",
    "\n",
    "    return _decorator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e527eb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "app = create_testing_app()\n",
    "\n",
    "\n",
    "# Basic check\n",
    "@app.consumes()\n",
    "def on_my_topic_1(msg: BaseModel) -> None:\n",
    "    pass\n",
    "\n",
    "\n",
    "assert app._consumers_store[\"my_topic_1_0\"] == (\n",
    "    on_my_topic_1,\n",
    "    json_decoder,\n",
    "    None,\n",
    "    app._kafka_brokers,\n",
    "    {},\n",
    "), app._consumers_store\n",
    "\n",
    "assert hasattr(app, \"on_my_topic_1\")\n",
    "\n",
    "\n",
    "# Check executor setting\n",
    "@app.consumes(executor=\"DynamicTaskExecutor\")\n",
    "def on_my_topic_12(msg: BaseModel) -> None:\n",
    "    pass\n",
    "\n",
    "\n",
    "assert app._consumers_store[\"my_topic_12_0\"] == (\n",
    "    on_my_topic_12,\n",
    "    json_decoder,\n",
    "    \"DynamicTaskExecutor\",\n",
    "    app._kafka_brokers,\n",
    "    {},\n",
    "), app._consumers_store[\"my_topic_12\"]\n",
    "\n",
    "assert hasattr(app, \"on_my_topic_12\")\n",
    "\n",
    "\n",
    "# Check topic setting\n",
    "@app.consumes(topic=\"test_topic_1\")\n",
    "def some_func_name(msg: BaseModel) -> None:\n",
    "    pass\n",
    "\n",
    "\n",
    "assert app._consumers_store[\"test_topic_1_0\"] == (\n",
    "    some_func_name,\n",
    "    json_decoder,\n",
    "    None,\n",
    "    app._kafka_brokers,\n",
    "    {},\n",
    "), app._consumers_store\n",
    "\n",
    "\n",
    "# Check prefix change\n",
    "@app.consumes(prefix=\"for_\")\n",
    "def for_test_topic_3(msg: BaseModel) -> None:\n",
    "    pass\n",
    "\n",
    "\n",
    "assert app._consumers_store[\"test_topic_3_0\"] == (\n",
    "    for_test_topic_3,\n",
    "    json_decoder,\n",
    "    None,\n",
    "    app._kafka_brokers,\n",
    "    {},\n",
    "), app._consumers_store\n",
    "\n",
    "assert hasattr(app, \"for_test_topic_3\")\n",
    "\n",
    "# Check passing of kwargs\n",
    "kwargs = {\"arg1\": \"val1\", \"arg2\": 2}\n",
    "\n",
    "\n",
    "@app.consumes(topic=\"test_topic\", **kwargs)\n",
    "def for_test_kwargs(msg: BaseModel):\n",
    "    pass\n",
    "\n",
    "\n",
    "assert app._consumers_store[\"test_topic_0\"] == (\n",
    "    for_test_kwargs,\n",
    "    json_decoder,\n",
    "    None,\n",
    "    app._kafka_brokers,\n",
    "    kwargs,\n",
    "), app._consumers_store\n",
    "\n",
    "assert hasattr(app, \"for_test_kwargs\")\n",
    "\n",
    "# Check description setting\n",
    "@app.consumes(description = \"Some generic description\")\n",
    "def on_test_topic_description(msg: BaseModel) -> None:\n",
    "    pass\n",
    "\n",
    "\n",
    "assert app._consumers_store[\"test_topic_description_0\"][0].description == \"Some generic description\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2fb7fb1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ok\n"
     ]
    }
   ],
   "source": [
    "# check broker overriding\n",
    "@app.consumes(\n",
    "    brokers=dict(\n",
    "        localhost=[\n",
    "            dict(url=\"localhost\", port=\"9092\"),\n",
    "        ]\n",
    "    ),\n",
    ")\n",
    "def on_my_topic_12345(msg: BaseModel) -> None:\n",
    "    pass\n",
    "\n",
    "\n",
    "expected = (\n",
    "    on_my_topic_12345,\n",
    "    json_decoder,\n",
    "    None,\n",
    "    KafkaBrokers(\n",
    "        brokers={\n",
    "            \"localhost\": [\n",
    "                KafkaBroker(\n",
    "                    url=\"localhost\",\n",
    "                    description=\"Kafka broker\",\n",
    "                    port=\"9092\",\n",
    "                    protocol=\"kafka\",\n",
    "                    security=None,\n",
    "                ),\n",
    "            ]\n",
    "        }\n",
    "    ),\n",
    "    {},\n",
    ")\n",
    "\n",
    "actual = app._consumers_store[\"my_topic_12345_0\"]\n",
    "\n",
    "assert actual == expected, f\"{actual}!={expected}\"\n",
    "\n",
    "assert hasattr(app, \"on_my_topic_12345\")\n",
    "\n",
    "with pytest.raises(ValueError) as e:\n",
    "\n",
    "    @app.consumes(\n",
    "        brokers=dict(\n",
    "            not_localhost=[\n",
    "                dict(url=\"localhost\", port=\"9092\"),\n",
    "                dict(url=\"localhost\", port=\"9093\"),\n",
    "            ]\n",
    "        ),\n",
    "    )\n",
    "    def on_my_topic_12345(msg: BaseModel) -> None:\n",
    "        pass\n",
    "\n",
    "\n",
    "assert (\n",
    "    e.value.args[0]\n",
    "    == \"To override application default brokers, you must define all of the broker options. Default defined: {'localhost'}, override defined: {'not_localhost'}\"\n",
    ")\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c01256d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _get_encoder_fn(encoder: str) -> Callable[[BaseModel], bytes]:\n",
    "    \"\"\"\n",
    "    Imports and returns encoder function based on input\n",
    "    \"\"\"\n",
    "    if encoder == \"json\":\n",
    "        from fastkafka._components.encoder.json import json_encoder\n",
    "\n",
    "        return json_encoder\n",
    "    elif encoder == \"avro\":\n",
    "        try:\n",
    "            from fastkafka._components.encoder.avro import avro_encoder\n",
    "        except ModuleNotFoundError:\n",
    "            raise ModuleNotFoundError(\n",
    "                \"Unable to import avro packages. Please install FastKafka using the command 'fastkafka[avro]'\"\n",
    "            )\n",
    "        return avro_encoder\n",
    "    else:\n",
    "        raise ValueError(f\"Unknown encoder - {encoder}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "804434e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "actual = _get_encoder_fn(\"json\")\n",
    "assert actual == json_encoder\n",
    "\n",
    "actual = _get_encoder_fn(\"avro\")\n",
    "assert actual == avro_encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e269659",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "@delegates(fastkafka._aiokafka_imports.AIOKafkaProducer)\n",
    "def produces(\n",
    "    self: FastKafka,\n",
    "    topic: Optional[str] = None,\n",
    "    encoder: Union[str, Callable[[BaseModel], bytes]] = \"json\",\n",
    "    *,\n",
    "    prefix: str = \"to_\",\n",
    "    brokers: Optional[Union[Dict[str, Any], KafkaBrokers]] = None,\n",
    "    description: Optional[str] = None,\n",
    "    **kwargs: Dict[str, Any],\n",
    ") -> Callable[[ProduceCallable], ProduceCallable]:\n",
    "    \"\"\"Decorator registering the callback called when delivery report for a produced message is received\n",
    "\n",
    "    This function decorator is also responsible for registering topics for AsyncAPI specificiation and documentation.\n",
    "\n",
    "    Args:\n",
    "        topic: Kafka topic that the producer will send returned values from\n",
    "            the decorated function to, default: None- If the topic is not\n",
    "            specified, topic name will be inferred from the decorated function\n",
    "            name by stripping the defined prefix.\n",
    "        encoder: Encoder to use to encode messages before sending it to topic,\n",
    "                default: json - By default, it uses json encoder to convert\n",
    "                pydantic basemodel to json string and then encodes the string to bytes\n",
    "                using 'utf-8' encoding. It also accepts custom encoder function.\n",
    "        prefix: Prefix stripped from the decorated function to define a topic\n",
    "            name if the topic argument is not passed, default: \"to_\". If the\n",
    "            decorated function name is not prefixed with the defined prefix\n",
    "            and topic argument is not passed, then this method will throw ValueError\n",
    "        brokers: Optional argument specifying multiple broker clusters for consuming\n",
    "            messages from different Kafka clusters in FastKafka.\n",
    "        description: Optional description of the producing function async docs.\n",
    "                If not provided, producing function __doc__ attr will be used.\n",
    "\n",
    "    Returns:\n",
    "        A function returning the same function\n",
    "\n",
    "    Raises:\n",
    "        ValueError: when needed\n",
    "    \"\"\"\n",
    "\n",
    "    def _decorator(\n",
    "        to_topic: ProduceCallable,\n",
    "        topic: Optional[str] = topic,\n",
    "        brokers: Optional[Union[Dict[str, Any], KafkaBrokers]] = brokers,\n",
    "        description: Optional[str] = description,\n",
    "        kwargs: Dict[str, Any] = kwargs,\n",
    "    ) -> ProduceCallable:\n",
    "        topic_resolved: str = (\n",
    "            _get_topic_name(topic_callable=to_topic, prefix=prefix)\n",
    "            if topic is None\n",
    "            else topic\n",
    "        )\n",
    "\n",
    "        topic_key = _resolve_key(topic_resolved, self._producers_store)\n",
    "\n",
    "        prepared_broker = _prepare_and_check_brokers(self, brokers)\n",
    "        if prepared_broker is not None:\n",
    "            self._override_brokers.append(prepared_broker.brokers)  # type: ignore\n",
    "        else:\n",
    "            prepared_broker = self._kafka_brokers\n",
    "\n",
    "        if description is not None:\n",
    "            setattr(to_topic, \"description\", description)\n",
    "\n",
    "        self._producers_store[topic_key] = (\n",
    "            to_topic,\n",
    "            None,\n",
    "            prepared_broker,\n",
    "            kwargs,\n",
    "        )\n",
    "        encoder_fn = _get_encoder_fn(encoder) if isinstance(encoder, str) else encoder\n",
    "        decorated = producer_decorator(\n",
    "            self._producers_store,\n",
    "            to_topic,\n",
    "            topic_key,\n",
    "            encoder_fn=encoder_fn,\n",
    "        )\n",
    "        setattr(self, to_topic.__name__, decorated)\n",
    "        return decorated\n",
    "\n",
    "    return _decorator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acf7cede",
   "metadata": {},
   "outputs": [],
   "source": [
    "app = create_testing_app()\n",
    "\n",
    "\n",
    "# Basic check\n",
    "async def to_my_topic_1(msg: BaseModel) -> None:\n",
    "    pass\n",
    "\n",
    "\n",
    "# Must be done without sugar to keep the original function reference\n",
    "check_func = to_my_topic_1\n",
    "to_my_topic_1 = app.produces()(to_my_topic_1)\n",
    "\n",
    "assert app._producers_store[\"my_topic_1_0\"] == (\n",
    "    check_func,\n",
    "    None,\n",
    "    app._kafka_brokers,\n",
    "    {},\n",
    "), f\"{app._producers_store}, {to_my_topic_1}\"\n",
    "\n",
    "assert hasattr(app, \"to_my_topic_1\")\n",
    "\n",
    "\n",
    "# Check topic setting\n",
    "async def some_func_name(msg: BaseModel) -> None:\n",
    "    pass\n",
    "\n",
    "\n",
    "check_func = some_func_name\n",
    "some_func_name = app.produces(topic=\"test_topic_2\")(some_func_name)\n",
    "\n",
    "assert app._producers_store[\"test_topic_2_0\"] == (\n",
    "    check_func,\n",
    "    None,\n",
    "    app._kafka_brokers,\n",
    "    {},\n",
    "), app._producers_store\n",
    "\n",
    "assert hasattr(app, \"some_func_name\")\n",
    "\n",
    "\n",
    "# Check prefix change\n",
    "async def for_test_topic_3(msg: BaseModel) -> None:\n",
    "    pass\n",
    "\n",
    "\n",
    "check_func = for_test_topic_3\n",
    "some_func_name = app.produces(prefix=\"for_\")(for_test_topic_3)\n",
    "\n",
    "assert app._producers_store[\"test_topic_3_0\"] == (\n",
    "    check_func,\n",
    "    None,\n",
    "    app._kafka_brokers,\n",
    "    {},\n",
    "), app._producers_store\n",
    "\n",
    "# Check passing of kwargs\n",
    "kwargs = {\"arg1\": \"val1\", \"arg2\": 2}\n",
    "\n",
    "assert hasattr(app, \"for_test_topic_3\")\n",
    "\n",
    "\n",
    "async def for_test_kwargs(msg: BaseModel):\n",
    "    pass\n",
    "\n",
    "\n",
    "check_func = for_test_kwargs\n",
    "for_test_kwargs = app.produces(topic=\"test_topic_0\", **kwargs)(for_test_kwargs)\n",
    "\n",
    "assert app._producers_store[\"test_topic_0_0\"] == (\n",
    "    check_func,\n",
    "    None,\n",
    "    app._kafka_brokers,\n",
    "    kwargs,\n",
    "), app._producers_store\n",
    "\n",
    "assert hasattr(app, \"for_test_kwargs\")\n",
    "\n",
    "# Check description setting\n",
    "async def to_test_topic_description(msg: BaseModel) -> None:\n",
    "    pass\n",
    "\n",
    "some_func_name = app.produces(description=\"Some generic producer\")(to_test_topic_description)\n",
    "\n",
    "assert app._producers_store[\"test_topic_description_0\"][0].description == \"Some generic producer\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35f1d574",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def to_test_topic_broker_override(msg: BaseModel) -> None:\n",
    "    pass\n",
    "\n",
    "\n",
    "check_func = to_test_topic_broker_override\n",
    "some_func_name = app.produces(\n",
    "    brokers=dict(\n",
    "        localhost=[\n",
    "            dict(url=\"localhost\", port=9092),\n",
    "        ]\n",
    "    ),\n",
    ")(to_test_topic_broker_override)\n",
    "\n",
    "assert app._producers_store[\"test_topic_broker_override_0\"] == (\n",
    "    check_func,\n",
    "    None,\n",
    "    KafkaBrokers(\n",
    "        brokers={\n",
    "            \"localhost\": [\n",
    "                KafkaBroker(\n",
    "                    url=\"localhost\",\n",
    "                    description=\"Kafka broker\",\n",
    "                    port=9092,\n",
    "                    protocol=\"kafka\",\n",
    "                    security=None,\n",
    "                ),\n",
    "            ]\n",
    "        }\n",
    "    ),\n",
    "    {},\n",
    "), app._producers_store[\"test_topic_broker_override_0\"]\n",
    "\n",
    "\n",
    "async def to_test_topic_broker_wrong(msg: BaseModel) -> None:\n",
    "    pass\n",
    "\n",
    "\n",
    "with pytest.raises(ValueError) as e:\n",
    "    check_func = to_test_topic_broker_wrong\n",
    "    some_func_name = app.produces(\n",
    "        brokers=dict(\n",
    "            not_localhost=[\n",
    "                dict(url=\"localhost\", port=9092),\n",
    "                dict(url=\"localhost\", port=9093),\n",
    "            ]\n",
    "        ),\n",
    "    )(to_test_topic_broker_wrong)\n",
    "\n",
    "assert (\n",
    "    e.value.args[0]\n",
    "    == \"To override application default brokers, you must define all of the broker options. Default defined: {'localhost'}, override defined: {'not_localhost'}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d507987",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "def get_topics(self: FastKafka) -> Iterable[str]:\n",
    "    \"\"\"\n",
    "    Get all topics for both producing and consuming.\n",
    "\n",
    "    Returns:\n",
    "        A set of topics for both producing and consuming.\n",
    "    \"\"\"\n",
    "    produce_topics = set([remove_suffix(topic) for topic in self._producers_store])\n",
    "    consume_topics = set([remove_suffix(topic) for topic in self._consumers_store])\n",
    "    return consume_topics.union(produce_topics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac053363",
   "metadata": {},
   "outputs": [],
   "source": [
    "app = create_testing_app()\n",
    "\n",
    "\n",
    "@app.produces()\n",
    "async def to_topic_1() -> BaseModel:\n",
    "    pass\n",
    "\n",
    "\n",
    "@app.consumes()\n",
    "def on_topic_2(msg: BaseModel):\n",
    "    pass\n",
    "\n",
    "\n",
    "assert app.get_topics() == set([\"topic_1\", \"topic_2\"]), f\"{app.get_topics()=}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4744bd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "def run_in_background(\n",
    "    self: FastKafka,\n",
    ") -> Callable[\n",
    "    [Callable[..., Coroutine[Any, Any, Any]]], Callable[..., Coroutine[Any, Any, Any]]\n",
    "]:\n",
    "    \"\"\"\n",
    "    Decorator to schedule a task to be run in the background.\n",
    "\n",
    "    This decorator is used to schedule a task to be run in the background when the app's `_on_startup` event is triggered.\n",
    "\n",
    "    Returns:\n",
    "        Callable[None, None]: A decorator function that takes a background task as an input and stores it to be run in the backround.\n",
    "    \"\"\"\n",
    "\n",
    "    def _decorator(\n",
    "        bg_task: Callable[..., Coroutine[Any, Any, Any]]\n",
    "    ) -> Callable[..., Coroutine[Any, Any, Any]]:\n",
    "        \"\"\"\n",
    "        Store the background task.\n",
    "\n",
    "        Args:\n",
    "            bg_task (Callable[[], None]): The background task to be run asynchronously.\n",
    "\n",
    "        Returns:\n",
    "            Callable[[], None]: Original background task.\n",
    "        \"\"\"\n",
    "        logger.info(\n",
    "            f\"run_in_background() : Adding function '{bg_task.__name__}' as background task\"\n",
    "        )\n",
    "        self._scheduled_bg_tasks.append(bg_task)\n",
    "\n",
    "        return bg_task\n",
    "\n",
    "    return _decorator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16917ec0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: run_in_background() : Adding function 'async_background_job' as background task\n"
     ]
    }
   ],
   "source": [
    "# Check if the background job is getting registered\n",
    "\n",
    "app = create_testing_app()\n",
    "\n",
    "\n",
    "@app.run_in_background()\n",
    "async def async_background_job():\n",
    "    \"\"\"Async background job\"\"\"\n",
    "    pass\n",
    "\n",
    "\n",
    "assert app._scheduled_bg_tasks[0] == async_background_job, app._scheduled_bg_tasks[0]\n",
    "assert app._scheduled_bg_tasks.__len__() == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e107bbeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyInfo(BaseModel):\n",
    "    mobile: str = Field(..., example=\"+385987654321\")\n",
    "    name: str = Field(..., example=\"James Bond\")\n",
    "\n",
    "\n",
    "class MyMsgUrl(BaseModel):\n",
    "    info: MyInfo = Field(..., example=dict(mobile=\"+385987654321\", name=\"James Bond\"))\n",
    "    url: HttpUrl = Field(..., example=\"https://sis.gov.uk/agents/007\")\n",
    "\n",
    "\n",
    "class MyMsgEmail(BaseModel):\n",
    "    msg_url: MyMsgUrl = Field(\n",
    "        ...,\n",
    "        example=dict(\n",
    "            info=dict(mobile=\"+385987654321\", name=\"James Bond\"),\n",
    "            url=\"https://sis.gov.uk/agents/007\",\n",
    "        ),\n",
    "    )\n",
    "    email: EmailStr = Field(..., example=\"agent-007@sis.gov.uk\")\n",
    "\n",
    "\n",
    "def setup_testing_app(bootstrap_servers=None, override_bootstrap_servers=None):\n",
    "    app = create_testing_app(bootstrap_servers=bootstrap_servers)\n",
    "\n",
    "    host, port = None, None\n",
    "    if override_bootstrap_servers is not None:\n",
    "        host, port = override_bootstrap_servers.split(\":\")\n",
    "\n",
    "    override_broker = {\n",
    "        \"localhost\": {\n",
    "            \"url\": host if host is not None else \"localhost\",\n",
    "            \"name\": \"development\",\n",
    "            \"description\": \"Local (dev) Kafka broker\",\n",
    "            \"port\": port if port is not None else \"9092\",\n",
    "        }\n",
    "    }\n",
    "\n",
    "    @app.consumes(\"my_topic_1\", description=\"Consumer description\")\n",
    "    def on_my_topic_one(msg: MyMsgUrl) -> None:\n",
    "        logger.debug(f\"on_my_topic_one(msg={msg},)\")\n",
    "\n",
    "    @app.consumes(topic=\"my_topic_1\", brokers=override_broker)\n",
    "    async def on_my_topic_1(msg: MyMsgEmail) -> None:\n",
    "        logger.debug(f\"on_my_topic_2(msg={msg},)\")\n",
    "\n",
    "    with pytest.raises(ValueError) as e:\n",
    "\n",
    "        @app.consumes()\n",
    "        def my_topic_3(msg: MyMsgEmail) -> None:\n",
    "            raise NotImplemented\n",
    "\n",
    "    @app.produces(description=\"Producer description\")\n",
    "    async def to_my_topic_3(url: str) -> MyMsgUrl:\n",
    "        logger.debug(f\"on_my_topic_3(msg={url}\")\n",
    "        return MyMsgUrl(info=MyInfo(\"+3851987654321\", \"Sean Connery\"), url=url)\n",
    "\n",
    "    @app.produces()\n",
    "    async def to_my_topic_4(msg: MyMsgEmail) -> MyMsgEmail:\n",
    "        logger.debug(f\"on_my_topic_4(msg={msg}\")\n",
    "        return msg\n",
    "\n",
    "    @app.produces(topic=\"my_topic_4\", brokers=override_broker)\n",
    "    async def to_my_topic_4_2(url: str) -> MyMsgUrl:\n",
    "        logger.debug(f\"on_my_topic_5(msg={url}\")\n",
    "        return MyMsgUrl(info=MyInfo(\"+3859123456789\", \"John Wayne\"), url=url)\n",
    "\n",
    "    @app.run_in_background()\n",
    "    async def long_bg_job():\n",
    "        logger.debug(f\"long_bg_job()\")\n",
    "        await asyncio.sleep(100)\n",
    "\n",
    "    return app"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a945425",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: run_in_background() : Adding function 'long_bg_job' as background task\n",
      "app._kafka_service_info=title='' version='' description='' contact=ContactInfo(name='Author', url=Url('https://www.google.com/'), email='noreply@gmail.com')\n",
      "app._kafka_brokers=brokers={'localhost': KafkaBroker(url='localhost', description='Local (dev) Kafka broker', port='9092', protocol='kafka', security=None)}\n"
     ]
    }
   ],
   "source": [
    "app = setup_testing_app()\n",
    "\n",
    "assert set(app._consumers_store.keys()) == set([\"my_topic_1_0\", \"my_topic_1_1\"])\n",
    "assert set(app._producers_store.keys()) == set(\n",
    "    [\"my_topic_3_0\", \"my_topic_4_0\", \"my_topic_4_1\"]\n",
    ")\n",
    "\n",
    "print(f\"app._kafka_service_info={app._kafka_service_info}\")\n",
    "print(f\"app._kafka_brokers={app._kafka_brokers}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb2ea338",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "def _populate_consumers(\n",
    "    self: FastKafka,\n",
    "    is_shutting_down_f: Callable[[], bool],\n",
    ") -> None:\n",
    "    default_config: Dict[str, Any] = filter_using_signature(\n",
    "        fastkafka._aiokafka_imports.AIOKafkaConsumer, **self._kafka_config\n",
    "    )\n",
    "\n",
    "    bootstrap_server = self._kafka_config[\"bootstrap_servers_id\"]\n",
    "\n",
    "    self._kafka_consumer_tasks = [\n",
    "        asyncio.create_task(\n",
    "            aiokafka_consumer_loop(\n",
    "                topic=\"_\".join(topic.split(\"_\")[:-1]),\n",
    "                decoder_fn=decoder_fn,\n",
    "                callback=consumer,\n",
    "                msg_type=signature(consumer).parameters[\"msg\"].annotation,\n",
    "                is_shutting_down_f=is_shutting_down_f,\n",
    "                executor=executor,\n",
    "                **{\n",
    "                    **default_config,\n",
    "                    **override_config,\n",
    "                    **{\n",
    "                        \"bootstrap_servers\": _get_broker_addr_list(\n",
    "                            kafka_brokers.brokers[bootstrap_server]\n",
    "                            if kafka_brokers is not None\n",
    "                            else self._kafka_brokers.brokers[bootstrap_server]\n",
    "                        )\n",
    "                    },\n",
    "                },\n",
    "            )\n",
    "        )\n",
    "        for topic, (\n",
    "            consumer,\n",
    "            decoder_fn,\n",
    "            executor,\n",
    "            kafka_brokers,\n",
    "            override_config,\n",
    "        ) in self._consumers_store.items()\n",
    "    ]\n",
    "\n",
    "\n",
    "@patch\n",
    "async def _shutdown_consumers(\n",
    "    self: FastKafka,\n",
    ") -> None:\n",
    "    if self._kafka_consumer_tasks:\n",
    "        await asyncio.wait(self._kafka_consumer_tasks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df10b5e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: But not exported to PATH, exporting...\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._components.test_dependencies: But not exported to PATH, exporting...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:15992\n",
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "Port 2181 is already in use\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: zookeeper startup failed, generating a new port and retrying...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: zookeeper new port=48943\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:15993\n",
      "[INFO] __main__: run_in_background() : Adding function 'long_bg_job' as background task\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': '127.0.0.1:15992'}\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': '127.0.0.1:15993'}\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'my_topic_1'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'my_topic_1'}\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'my_topic_1'})\n",
      "[INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'my_topic_1'}\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[WARNING] aiokafka.cluster: Topic my_topic_1 is not available during auto-create initialization\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'my_topic_1': 0}. \n",
      "[WARNING] aiokafka.cluster: Topic my_topic_1 is not available during auto-create initialization\n",
      "[INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'my_topic_1': 0}. \n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 1886...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 1886 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 1516...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 1516 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 1088...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 1088 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 718...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 718 terminated.\n"
     ]
    }
   ],
   "source": [
    "async with ApacheKafkaBroker(listener_port=15992) as bootstrap_server:\n",
    "    async with ApacheKafkaBroker(listener_port=15993) as override_bootstrap_server:\n",
    "        app = setup_testing_app(\n",
    "            bootstrap_servers=bootstrap_server,\n",
    "            override_bootstrap_servers=override_bootstrap_server,\n",
    "        )\n",
    "        app._populate_consumers(is_shutting_down_f=true_after(1))\n",
    "        assert len(app._kafka_consumer_tasks) == 2\n",
    "\n",
    "        await app._shutdown_consumers()\n",
    "\n",
    "        assert all([t.done() for t in app._kafka_consumer_tasks])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d6ee2b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "# TODO: Add passing of vars\n",
    "async def _create_producer(  # type: ignore\n",
    "    *,\n",
    "    callback: ProduceCallable,\n",
    "    default_config: Dict[str, Any],\n",
    "    override_config: Dict[str, Any],\n",
    "    bootstrap_servers: Union[str, List[str]],\n",
    "    producers_list: List[fastkafka._aiokafka_imports.AIOKafkaProducer],\n",
    ") -> fastkafka._aiokafka_imports.AIOKafkaProducer:\n",
    "    \"\"\"Creates a producer\n",
    "\n",
    "    Args:\n",
    "        callback: A callback function that is called when the producer is ready.\n",
    "        producer: An existing producer to use.\n",
    "        default_config: A dictionary of default configuration values.\n",
    "        override_config: A dictionary of configuration values to override.\n",
    "        bootstrap_servers: Bootstrap servers to connect the producer to.\n",
    "        producers_list: A list of producers to add the new producer to.\n",
    "\n",
    "    Returns:\n",
    "        A producer.\n",
    "    \"\"\"\n",
    "\n",
    "    config = {\n",
    "        **filter_using_signature(\n",
    "            fastkafka._aiokafka_imports.AIOKafkaProducer, **default_config\n",
    "        ),\n",
    "        **filter_using_signature(\n",
    "            fastkafka._aiokafka_imports.AIOKafkaProducer, **override_config\n",
    "        ),\n",
    "        **{\"bootstrap_servers\": bootstrap_servers},\n",
    "    }\n",
    "\n",
    "    producer = fastkafka._aiokafka_imports.AIOKafkaProducer(**config)\n",
    "    logger.info(\n",
    "        f\"_create_producer() : created producer using the config: '{sanitize_kafka_config(**config)}'\"\n",
    "    )\n",
    "\n",
    "    await producer.start()\n",
    "\n",
    "    producers_list.append(producer)\n",
    "\n",
    "    return producer\n",
    "\n",
    "\n",
    "@patch\n",
    "async def _populate_producers(self: FastKafka) -> None:\n",
    "    \"\"\"Populates the producers for the FastKafka instance.\n",
    "\n",
    "    Args:\n",
    "        self: The FastKafka instance.\n",
    "\n",
    "    Returns:\n",
    "        None.\n",
    "\n",
    "    Raises:\n",
    "        None.\n",
    "    \"\"\"\n",
    "    default_config: Dict[str, Any] = self._kafka_config\n",
    "    bootstrap_server = default_config[\"bootstrap_servers_id\"]\n",
    "\n",
    "    self._producers_list = []\n",
    "    self._producers_store.update(\n",
    "        {\n",
    "            topic: (\n",
    "                callback,\n",
    "                await _create_producer(\n",
    "                    callback=callback,\n",
    "                    default_config=default_config,\n",
    "                    override_config=override_config,\n",
    "                    bootstrap_servers=_get_broker_addr_list(\n",
    "                        kafka_brokers.brokers[bootstrap_server]\n",
    "                        if kafka_brokers is not None\n",
    "                        else self._kafka_brokers.brokers[bootstrap_server]\n",
    "                    ),\n",
    "                    producers_list=self._producers_list,\n",
    "                ),\n",
    "                kafka_brokers,\n",
    "                override_config,\n",
    "            )\n",
    "            for topic, (\n",
    "                callback,\n",
    "                _,\n",
    "                kafka_brokers,\n",
    "                override_config,\n",
    "            ) in self._producers_store.items()\n",
    "        }\n",
    "    )\n",
    "\n",
    "\n",
    "@patch\n",
    "async def _shutdown_producers(self: FastKafka) -> None:\n",
    "    [await producer.stop() for producer in self._producers_list[::-1]]\n",
    "    # Remove references to stale producers\n",
    "    self._producers_list = []\n",
    "    self._producers_store.update(\n",
    "        {\n",
    "            topic: (\n",
    "                callback,\n",
    "                None,\n",
    "                kafka_brokers,\n",
    "                override_config,\n",
    "            )\n",
    "            for topic, (\n",
    "                callback,\n",
    "                _,\n",
    "                kafka_brokers,\n",
    "                override_config,\n",
    "            ) in self._producers_store.items()\n",
    "        }\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0546037d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:15992\n",
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "Port 2181 is already in use\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: zookeeper startup failed, generating a new port and retrying...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: zookeeper new port=34285\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:15993\n",
      "[INFO] __main__: run_in_background() : Adding function 'long_bg_job' as background task\n",
      "{'my_topic_3_0': (<function setup_testing_app.<locals>.to_my_topic_3>, None, KafkaBrokers(brokers={'localhost': KafkaBroker(url='127.0.0.1', description='Local (dev) Kafka broker', port='15992', protocol='kafka', security=None)}), {}), 'my_topic_4_0': (<function setup_testing_app.<locals>.to_my_topic_4>, None, KafkaBrokers(brokers={'localhost': KafkaBroker(url='127.0.0.1', description='Local (dev) Kafka broker', port='15992', protocol='kafka', security=None)}), {}), 'my_topic_4_1': (<function setup_testing_app.<locals>.to_my_topic_4_2>, None, KafkaBrokers(brokers={'localhost': KafkaBroker(url='127.0.0.1', description='Local (dev) Kafka broker', port='15993', protocol='kafka', security=None)}), {})}\n",
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': '127.0.0.1:15992'}'\n",
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': '127.0.0.1:15992'}'\n",
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': '127.0.0.1:15993'}'\n",
      "{'my_topic_3_0': (<function setup_testing_app.<locals>.to_my_topic_3>, <aiokafka.producer.producer.AIOKafkaProducer object>, KafkaBrokers(brokers={'localhost': KafkaBroker(url='127.0.0.1', description='Local (dev) Kafka broker', port='15992', protocol='kafka', security=None)}), {}), 'my_topic_4_0': (<function setup_testing_app.<locals>.to_my_topic_4>, <aiokafka.producer.producer.AIOKafkaProducer object>, KafkaBrokers(brokers={'localhost': KafkaBroker(url='127.0.0.1', description='Local (dev) Kafka broker', port='15992', protocol='kafka', security=None)}), {}), 'my_topic_4_1': (<function setup_testing_app.<locals>.to_my_topic_4_2>, <aiokafka.producer.producer.AIOKafkaProducer object>, KafkaBrokers(brokers={'localhost': KafkaBroker(url='127.0.0.1', description='Local (dev) Kafka broker', port='15993', protocol='kafka', security=None)}), {})}\n",
      "[<aiokafka.producer.producer.AIOKafkaProducer object>, <aiokafka.producer.producer.AIOKafkaProducer object>, <aiokafka.producer.producer.AIOKafkaProducer object>]\n",
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': '127.0.0.1:15992'}'\n",
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': '127.0.0.1:15992'}'\n",
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': '127.0.0.1:15993'}'\n",
      "[<aiokafka.producer.producer.AIOKafkaProducer object>, <aiokafka.producer.producer.AIOKafkaProducer object>, <aiokafka.producer.producer.AIOKafkaProducer object>]\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 3518...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 3518 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 3149...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 3149 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 2720...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 2720 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 2350...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 2350 terminated.\n"
     ]
    }
   ],
   "source": [
    "async with ApacheKafkaBroker(listener_port=15992) as bootstrap_server:\n",
    "    async with ApacheKafkaBroker(listener_port=15993) as override_bootstrap_server:\n",
    "        app = setup_testing_app(\n",
    "            bootstrap_servers=bootstrap_server,\n",
    "            override_bootstrap_servers=override_bootstrap_server,\n",
    "        )\n",
    "        print(app._producers_store)\n",
    "        await app._populate_producers()\n",
    "        print(app._producers_store)\n",
    "        assert len(app._producers_list) == 3\n",
    "        print(app._producers_list)\n",
    "        await app._shutdown_producers()\n",
    "\n",
    "        # One more time for reentrancy\n",
    "        await app._populate_producers()\n",
    "        assert len(app._producers_list) == 3\n",
    "        print(app._producers_list)\n",
    "        await app._shutdown_producers()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b15e020",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "async def _populate_bg_tasks(\n",
    "    self: FastKafka,\n",
    ") -> None:\n",
    "    def _start_bg_task(task: Callable[..., Coroutine[Any, Any, Any]]) -> asyncio.Task:\n",
    "        logger.info(\n",
    "            f\"_populate_bg_tasks() : Starting background task '{task.__name__}'\"\n",
    "        )\n",
    "        return asyncio.create_task(task(), name=task.__name__)\n",
    "\n",
    "    self._running_bg_tasks = [_start_bg_task(task) for task in self._scheduled_bg_tasks]\n",
    "\n",
    "\n",
    "@patch\n",
    "async def _shutdown_bg_tasks(\n",
    "    self: FastKafka,\n",
    ") -> None:\n",
    "    for task in self._running_bg_tasks:\n",
    "        logger.info(\n",
    "            f\"_shutdown_bg_tasks() : Cancelling background task '{task.get_name()}'\"\n",
    "        )\n",
    "        task.cancel()\n",
    "\n",
    "    for task in self._running_bg_tasks:\n",
    "        logger.info(\n",
    "            f\"_shutdown_bg_tasks() : Waiting for background task '{task.get_name()}' to finish\"\n",
    "        )\n",
    "        try:\n",
    "            await task\n",
    "        except asyncio.CancelledError:\n",
    "            pass\n",
    "        logger.info(\n",
    "            f\"_shutdown_bg_tasks() : Execution finished for background task '{task.get_name()}'\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4c687d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:15992\n",
      "[INFO] __main__: run_in_background() : Adding function 'long_bg_job' as background task\n",
      "[INFO] __main__: run_in_background() : Adding function 'long_bg_job' as background task\n",
      "[INFO] __main__: _populate_bg_tasks() : Starting background task 'long_bg_job'\n",
      "[INFO] __main__: _populate_bg_tasks() : Starting background task 'long_bg_job'\n",
      "[INFO] __main__: _shutdown_bg_tasks() : Cancelling background task 'long_bg_job'\n",
      "[INFO] __main__: _shutdown_bg_tasks() : Cancelling background task 'long_bg_job'\n",
      "[INFO] __main__: _shutdown_bg_tasks() : Waiting for background task 'long_bg_job' to finish\n",
      "[INFO] __main__: _shutdown_bg_tasks() : Execution finished for background task 'long_bg_job'\n",
      "[INFO] __main__: _shutdown_bg_tasks() : Waiting for background task 'long_bg_job' to finish\n",
      "[INFO] __main__: _shutdown_bg_tasks() : Execution finished for background task 'long_bg_job'\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 4348...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 4348 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 3977...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 3977 terminated.\n"
     ]
    }
   ],
   "source": [
    "async with ApacheKafkaBroker(listener_port=15992) as bootstrap_server:\n",
    "    app = setup_testing_app(bootstrap_servers=bootstrap_server)\n",
    "\n",
    "    @app.run_in_background()\n",
    "    async def long_bg_job():\n",
    "        logger.debug(f\"new_long_bg_job()\")\n",
    "        await asyncio.sleep(100)\n",
    "\n",
    "    await app._populate_bg_tasks()\n",
    "    assert len(app._scheduled_bg_tasks) == 2\n",
    "    assert len(app._running_bg_tasks) == 2\n",
    "    await app._shutdown_bg_tasks()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13bd84f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "async def _start(self: FastKafka) -> None:\n",
    "    def is_shutting_down_f(self: FastKafka = self) -> bool:\n",
    "        return self._is_shutting_down\n",
    "\n",
    "    #     self.create_docs()\n",
    "    await self._populate_producers()\n",
    "    self._populate_consumers(is_shutting_down_f)\n",
    "    await self._populate_bg_tasks()\n",
    "\n",
    "    self._is_started = True\n",
    "\n",
    "\n",
    "@patch\n",
    "async def _stop(self: FastKafka) -> None:\n",
    "    self._is_shutting_down = True\n",
    "\n",
    "    await self._shutdown_bg_tasks()\n",
    "    await self._shutdown_consumers()\n",
    "    await self._shutdown_producers()\n",
    "\n",
    "    self._is_shutting_down = False\n",
    "    self._is_started = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18b199a8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:15992\n",
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': '127.0.0.1:15992'}'\n",
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': '127.0.0.1:15992'}'\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 5160...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 5160 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 4791...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 4791 terminated.\n"
     ]
    }
   ],
   "source": [
    "# Test app reentrancy\n",
    "\n",
    "async with ApacheKafkaBroker(listener_port=15992) as bootstrap_server:\n",
    "    with mock_AIOKafkaProducer_send() as mock:\n",
    "        app = create_testing_app(bootstrap_servers=bootstrap_server)\n",
    "\n",
    "        @app.produces()\n",
    "        async def to_my_test_topic(mobile: str, url: str) -> MyMsgUrl:\n",
    "            msg = MyMsgUrl(info=dict(mobile=mobile, name=\"James Bond\"), url=url)\n",
    "            return msg\n",
    "\n",
    "        try:\n",
    "            await app._start()\n",
    "            await app.to_my_test_topic(mobile=\"+385912345678\", url=\"https://www.vip.hr\")\n",
    "        finally:\n",
    "            await app._stop()\n",
    "\n",
    "        try:\n",
    "            await app._start()\n",
    "            await app.to_my_test_topic(mobile=\"+385987654321\", url=\"https://www.ht.hr\")\n",
    "        finally:\n",
    "            await app._stop()\n",
    "\n",
    "        mock.assert_has_calls(\n",
    "            [\n",
    "                unittest.mock.call(\n",
    "                    \"my_test_topic\",\n",
    "                    b'{\"info\":{\"mobile\":\"+385912345678\",\"name\":\"James Bond\"},\"url\":\"https://www.vip.hr/\"}',\n",
    "                    key=None,\n",
    "                ),\n",
    "                unittest.mock.call(\n",
    "                    \"my_test_topic\",\n",
    "                    b'{\"info\":{\"mobile\":\"+385987654321\",\"name\":\"James Bond\"},\"url\":\"https://www.ht.hr/\"}',\n",
    "                    key=None,\n",
    "                ),\n",
    "            ]\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff9cfce8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:15992\n",
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': '127.0.0.1:15992'}'\n",
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': '127.0.0.1:15992'}'\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 5975...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 5975 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 5604...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 5604 terminated.\n"
     ]
    }
   ],
   "source": [
    "# mock up send method offastkafka._aiokafka_imports.AIOKafkaProducer\n",
    "async with ApacheKafkaBroker(listener_port=15992) as bootstrap_server:\n",
    "    with mock_AIOKafkaProducer_send() as mock:\n",
    "        app = create_testing_app(bootstrap_servers=bootstrap_server)\n",
    "\n",
    "        @app.produces()\n",
    "        async def to_my_test_topic(mobile: str, url: str) -> MyMsgUrl:\n",
    "            msg = MyMsgUrl(info=dict(mobile=mobile, name=\"James Bond\"), url=url)\n",
    "            return msg\n",
    "\n",
    "        @app.produces()\n",
    "        async def to_my_test_topic_2(mobile: str, url: str) -> MyMsgUrl:\n",
    "            msg = MyMsgUrl(info=dict(mobile=mobile, name=\"James Bond\"), url=url)\n",
    "            return msg\n",
    "\n",
    "        try:\n",
    "            await app._start()\n",
    "            await to_my_test_topic(mobile=\"+385912345678\", url=\"https://www.vip.hr\")\n",
    "            await to_my_test_topic_2(mobile=\"+385987654321\", url=\"https://www.ht.hr\")\n",
    "        finally:\n",
    "            await app._stop()\n",
    "\n",
    "        mock.assert_has_calls(\n",
    "            [\n",
    "                unittest.mock.call(\n",
    "                    \"my_test_topic\",\n",
    "                    b'{\"info\":{\"mobile\":\"+385912345678\",\"name\":\"James Bond\"},\"url\":\"https://www.vip.hr/\"}',\n",
    "                    key=None,\n",
    "                ),\n",
    "                unittest.mock.call(\n",
    "                    \"my_test_topic_2\",\n",
    "                    b'{\"info\":{\"mobile\":\"+385987654321\",\"name\":\"James Bond\"},\"url\":\"https://www.ht.hr/\"}',\n",
    "                    key=None,\n",
    "                ),\n",
    "            ]\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b114fc4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:15992\n",
      "[INFO] __main__: run_in_background() : Adding function 'bg_task' as background task\n",
      "[INFO] __main__: run_in_background() : Adding function 'bg_task_second' as background task\n",
      "[INFO] __main__: _populate_bg_tasks() : Starting background task 'bg_task'\n",
      "[INFO] __main__: _populate_bg_tasks() : Starting background task 'bg_task_second'\n",
      "[INFO] __main__: _shutdown_bg_tasks() : Cancelling background task 'bg_task'\n",
      "[INFO] __main__: _shutdown_bg_tasks() : Cancelling background task 'bg_task_second'\n",
      "[INFO] __main__: _shutdown_bg_tasks() : Waiting for background task 'bg_task' to finish\n",
      "[INFO] __main__: _shutdown_bg_tasks() : Execution finished for background task 'bg_task'\n",
      "[INFO] __main__: _shutdown_bg_tasks() : Waiting for background task 'bg_task_second' to finish\n",
      "[INFO] __main__: _shutdown_bg_tasks() : Execution finished for background task 'bg_task_second'\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 6787...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 6787 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 6418...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 6418 terminated.\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "async with ApacheKafkaBroker(listener_port=15992) as bootstrap_server:\n",
    "    app = create_testing_app(bootstrap_servers=bootstrap_server)\n",
    "    fast_task = unittest.mock.Mock()\n",
    "    long_task = unittest.mock.Mock()\n",
    "\n",
    "    @app.run_in_background()\n",
    "    async def bg_task():\n",
    "        fast_task()\n",
    "        await asyncio.sleep(100)\n",
    "        long_task()\n",
    "\n",
    "    fast_task_second = unittest.mock.Mock()\n",
    "    long_task_second = unittest.mock.Mock()\n",
    "\n",
    "    @app.run_in_background()\n",
    "    async def bg_task_second():\n",
    "        fast_task_second()\n",
    "        await asyncio.sleep(100)\n",
    "        long_task_second()\n",
    "\n",
    "    try:\n",
    "        await app._start()\n",
    "        await asyncio.sleep(5)\n",
    "    finally:\n",
    "        await app._stop()\n",
    "\n",
    "    fast_task.assert_called()\n",
    "    long_task.assert_not_called()\n",
    "\n",
    "    fast_task_second.assert_called()\n",
    "    long_task_second.assert_not_called()\n",
    "\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c4bb0ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.apache_kafka_broker: ApacheKafkaBroker.start(): entering...\n",
      "[WARNING] fastkafka._testing.apache_kafka_broker: ApacheKafkaBroker.start(): (<_UnixSelectorEventLoop running=True closed=False debug=False>) is already running!\n",
      "[WARNING] fastkafka._testing.apache_kafka_broker: ApacheKafkaBroker.start(): calling nest_asyncio.apply()\n",
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:15992\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: <class 'fastkafka.testing.ApacheKafkaBroker'>.start(): returning 127.0.0.1:15992\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: ApacheKafkaBroker.start(): exited.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: ApacheKafkaBroker.stop(): entering...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 7600...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 7600 terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 7231...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 7231 terminated.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: ApacheKafkaBroker.stop(): exited.\n"
     ]
    }
   ],
   "source": [
    "# test lifespan hook\n",
    "\n",
    "global_dict = {}\n",
    "\n",
    "\n",
    "@asynccontextmanager\n",
    "async def lifespan(app: FastKafka):\n",
    "    try:\n",
    "        global_dict[\"set_var\"] = 123\n",
    "        global_dict[\"app\"] = app\n",
    "        yield\n",
    "    finally:\n",
    "        global_dict[\"set_var\"] = 321\n",
    "\n",
    "\n",
    "with ApacheKafkaBroker(apply_nest_asyncio=True, listener_port=15992) as bootstrap_servers:\n",
    "    host, port = bootstrap_servers.split(\":\")\n",
    "\n",
    "    kafka_app = FastKafka(\n",
    "        kafka_brokers={\n",
    "            \"localhost\": {\n",
    "                \"url\": host if host is not None else \"localhost\",\n",
    "                \"name\": \"development\",\n",
    "                \"description\": \"Local (dev) Kafka broker\",\n",
    "                \"port\": port if port is not None else 9092,\n",
    "            }\n",
    "        },\n",
    "        root_path=\"/tmp/000_FastKafka\",\n",
    "        lifespan=lifespan,\n",
    "    )\n",
    "\n",
    "    kafka_app.set_kafka_broker(kafka_broker_name=\"localhost\")\n",
    "\n",
    "    # Dict unchanged\n",
    "    assert global_dict == {}\n",
    "\n",
    "    async with kafka_app:\n",
    "        # Lifespan aenter triggered\n",
    "        assert global_dict[\"set_var\"] == 123\n",
    "        # Kafka app reference passed\n",
    "        assert global_dict[\"app\"] == kafka_app\n",
    "\n",
    "    # Lifespan aexit triggered\n",
    "    assert global_dict[\"set_var\"] == 321"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe047872",
   "metadata": {},
   "source": [
    "## Documentation generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "880411a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "def create_docs(self: FastKafka) -> None:\n",
    "    \"\"\"\n",
    "    Create the asyncapi documentation based on the configured consumers and producers.\n",
    "\n",
    "    This function exports the asyncapi specification based on the configured consumers\n",
    "    and producers in the FastKafka instance. It generates the asyncapi documentation by\n",
    "    extracting the topics and callbacks from the consumers and producers.\n",
    "\n",
    "    Note:\n",
    "        The asyncapi documentation is saved to the location specified by the `_asyncapi_path`\n",
    "        attribute of the FastKafka instance.\n",
    "    \"\"\"\n",
    "    (self._asyncapi_path / \"docs\").mkdir(exist_ok=True, parents=True)\n",
    "    (self._asyncapi_path / \"spec\").mkdir(exist_ok=True, parents=True)\n",
    "    export_async_spec(\n",
    "        consumers={\n",
    "            remove_suffix(topic) if topic.endswith(\"_0\") else topic: callback\n",
    "            for topic, (callback, _, _, _, _) in self._consumers_store.items()\n",
    "        },\n",
    "        producers={\n",
    "            remove_suffix(topic) if topic.endswith(\"_0\") else topic: callback\n",
    "            for topic, (callback, _, _, _) in self._producers_store.items()\n",
    "        },\n",
    "        kafka_brokers=self._kafka_brokers,\n",
    "        kafka_service_info=self._kafka_service_info,\n",
    "        asyncapi_path=self._asyncapi_path,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ec33cb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "expected = \"\"\"asyncapi: 2.5.0\n",
    "channels:\n",
    "  my_topic_1:\n",
    "    subscribe:\n",
    "      message:\n",
    "        $ref: '#/components/messages/MyMsgUrl'\n",
    "      description: \"Consumer description\"\n",
    "  my_topic_1_1:\n",
    "    subscribe:\n",
    "      message:\n",
    "        $ref: '#/components/messages/MyMsgEmail'\n",
    "  my_topic_3:\n",
    "    publish:\n",
    "      message:\n",
    "        $ref: '#/components/messages/MyMsgUrl'\n",
    "      description: \"Producer description\"\n",
    "  my_topic_4:\n",
    "    publish:\n",
    "      message:\n",
    "        $ref: '#/components/messages/MyMsgEmail'\n",
    "  my_topic_4_1:\n",
    "    publish:\n",
    "      message:\n",
    "        $ref: '#/components/messages/MyMsgUrl'\n",
    "components:\n",
    "  messages:\n",
    "    MyMsgEmail:\n",
    "      payload:\n",
    "        example:\n",
    "          email: agent-007@sis.gov.uk\n",
    "          msg_url:\n",
    "            info:\n",
    "              mobile: '+385987654321'\n",
    "              name: James Bond\n",
    "            url: https://sis.gov.uk/agents/007\n",
    "        properties:\n",
    "          email:\n",
    "            example: agent-007@sis.gov.uk\n",
    "            format: email\n",
    "            title: Email\n",
    "            type: string\n",
    "          msg_url:\n",
    "            allOf:\n",
    "            - $ref: '#/components/messages/MyMsgUrl'\n",
    "            example:\n",
    "              info:\n",
    "                mobile: '+385987654321'\n",
    "                name: James Bond\n",
    "              url: https://sis.gov.uk/agents/007\n",
    "        required:\n",
    "        - msg_url\n",
    "        - email\n",
    "        title: MyMsgEmail\n",
    "        type: object\n",
    "    MyMsgUrl:\n",
    "      payload:\n",
    "        example:\n",
    "          info:\n",
    "            mobile: '+385987654321'\n",
    "            name: James Bond\n",
    "          url: https://sis.gov.uk/agents/007\n",
    "        properties:\n",
    "          info:\n",
    "            allOf:\n",
    "            - $ref: '#/components/schemas/MyInfo'\n",
    "            example:\n",
    "              mobile: '+385987654321'\n",
    "              name: James Bond\n",
    "          url:\n",
    "            example: https://sis.gov.uk/agents/007\n",
    "            format: uri\n",
    "            maxLength: 2083\n",
    "            minLength: 1\n",
    "            title: Url\n",
    "            type: string\n",
    "        required:\n",
    "        - info\n",
    "        - url\n",
    "        title: MyMsgUrl\n",
    "        type: object\n",
    "  schemas:\n",
    "    MyInfo:\n",
    "      payload:\n",
    "        properties:\n",
    "          mobile:\n",
    "            example: '+385987654321'\n",
    "            title: Mobile\n",
    "            type: string\n",
    "          name:\n",
    "            example: James Bond\n",
    "            title: Name\n",
    "            type: string\n",
    "        required:\n",
    "        - mobile\n",
    "        - name\n",
    "        title: MyInfo\n",
    "        type: object\n",
    "  securitySchemes: {}\n",
    "info:\n",
    "  contact:\n",
    "    email: noreply@gmail.com\n",
    "    name: Author\n",
    "    url: https://www.google.com/\n",
    "  description: ''\n",
    "  title: ''\n",
    "  version: ''\n",
    "servers:\n",
    "  localhost:\n",
    "    description: Local (dev) Kafka broker\n",
    "    protocol: kafka\n",
    "    url: localhost\n",
    "    variables:\n",
    "      port:\n",
    "        default: '9092'\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c407913e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: run_in_background() : Adding function 'long_bg_job' as background task\n",
      "[INFO] fastkafka._components.asyncapi: Old async specifications at '/tmp/000_FastKafka/asyncapi/spec/asyncapi.yml' does not exist.\n",
      "[INFO] fastkafka._components.asyncapi: New async specifications generated at: '/tmp/000_FastKafka/asyncapi/spec/asyncapi.yml'\n",
      "[INFO] fastkafka._components.asyncapi: Async docs generated at '/tmp/000_FastKafka/asyncapi/docs'\n",
      "[INFO] fastkafka._components.asyncapi: Output of '$ npx -y -p @asyncapi/generator ag /tmp/000_FastKafka/asyncapi/spec/asyncapi.yml @asyncapi/html-template -o /tmp/000_FastKafka/asyncapi/docs --force-write'npm WARN deprecated har-validator@5.1.5: this library is no longer supported\n",
      "npm WARN deprecated debuglog@1.0.1: Package no longer supported. Contact Support at https://www.npmjs.com/support for more info.\n",
      "npm WARN deprecated uuid@3.4.0: Please upgrade  to version 7 or higher.  Older versions may use Math.random() in certain circumstances, which is known to be problematic.  See https://v8.dev/blog/math-random for details.\n",
      "npm WARN deprecated request@2.88.2: request has been deprecated, see https://github.com/request/request/issues/3142\n",
      "npm WARN deprecated readdir-scoped-modules@1.1.0: This functionality has been moved to @npmcli/fs\n",
      "npm WARN deprecated @npmcli/move-file@1.1.2: This functionality has been moved to @npmcli/fs\n",
      "npm WARN deprecated mkdirp@0.3.5: Legacy versions of mkdirp are no longer supported. Please update to mkdirp 1.x. (Note that the API surface has changed to use Promises in 1.x.)\n",
      "npm WARN deprecated mkdirp@0.3.5: Legacy versions of mkdirp are no longer supported. Please update to mkdirp 1.x. (Note that the API surface has changed to use Promises in 1.x.)\n",
      "\u001b[32m\n",
      "\n",
      "Done! ✨\u001b[0m\n",
      "\u001b[33mCheck out your shiny new generated files at \u001b[0m\u001b[35m/tmp/000_FastKafka/asyncapi/docs\u001b[0m\u001b[33m.\u001b[0m\n",
      "\n",
      "npm notice \n",
      "npm notice New minor version of npm available! 9.6.3 -> 9.8.1\n",
      "npm notice Changelog: <https://github.com/npm/cli/releases/tag/v9.8.1>\n",
      "npm notice Run `npm install -g npm@9.8.1` to update!\n",
      "npm notice \n",
      "\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "d1, d2 = None, None\n",
    "\n",
    "docs_path = Path(\"/tmp/000_FastKafka/asyncapi/spec/asyncapi.yml\")\n",
    "if docs_path.exists():\n",
    "    os.remove(docs_path)\n",
    "\n",
    "\n",
    "async def test_me():\n",
    "    global d1\n",
    "    global d2\n",
    "    app = setup_testing_app()\n",
    "    app.create_docs()\n",
    "    with open(docs_path) as specs:\n",
    "        d1 = yaml.safe_load(specs)\n",
    "        d2 = yaml.safe_load(expected)\n",
    "        assert d1 == d2, f\"{d1} != {d2}\"\n",
    "\n",
    "\n",
    "asyncio.run(test_me())\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cafbc67",
   "metadata": {},
   "source": [
    "## App mocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4cdce7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "class AwaitedMock:\n",
    "    \"\"\"\n",
    "    Class representing an awaited mock object.\n",
    "\n",
    "    Args:\n",
    "        o: The original object to be wrapped.\n",
    "    \"\"\"\n",
    "\n",
    "    @staticmethod\n",
    "    def _await_for(f: Callable[..., Any]) -> Callable[..., Any]:\n",
    "        @delegates(f)\n",
    "        async def inner(\n",
    "            *args: Any, f: Callable[..., Any] = f, timeout: int = 60, **kwargs: Any\n",
    "        ) -> Any:\n",
    "            \"\"\"\n",
    "            Decorator to await the execution of a function.\n",
    "\n",
    "            Args:\n",
    "                f: The function to be wrapped.\n",
    "\n",
    "            Returns:\n",
    "                The wrapped function.\n",
    "            \"\"\"\n",
    "            if inspect.iscoroutinefunction(f):\n",
    "                return await asyncio.wait_for(f(*args, **kwargs), timeout=timeout)\n",
    "            else:\n",
    "                t0 = datetime.now()\n",
    "                e: Optional[Exception] = None\n",
    "                while True:\n",
    "                    try:\n",
    "                        return f(*args, **kwargs)\n",
    "                    except Exception as _e:\n",
    "                        await asyncio.sleep(1)\n",
    "                        e = _e\n",
    "\n",
    "                    if datetime.now() - t0 > timedelta(seconds=timeout):\n",
    "                        break\n",
    "\n",
    "                raise e\n",
    "\n",
    "        return inner\n",
    "\n",
    "    def __init__(self, o: Any):\n",
    "        \"\"\"\n",
    "        Initializes an instance of AwaitedMock.\n",
    "\n",
    "        Args:\n",
    "            o: The original object to be wrapped.\n",
    "        \"\"\"\n",
    "        self._o = o\n",
    "\n",
    "        for name in o.__dir__():\n",
    "            if not name.startswith(\"_\"):\n",
    "                f = getattr(o, name)\n",
    "                if inspect.ismethod(f):\n",
    "                    setattr(self, name, self._await_for(f))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d95b332",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "def create_mocks(self: FastKafka) -> None:\n",
    "    \"\"\"Creates self.mocks as a named tuple mapping a new function obtained by calling the original functions and a mock\"\"\"\n",
    "    app_methods = [f for f, _, _, _, _ in self._consumers_store.values()] + [\n",
    "        f for f, _, _, _ in self._producers_store.values()\n",
    "    ]\n",
    "    self.AppMocks = namedtuple(  # type: ignore\n",
    "        f\"{self.__class__.__name__}Mocks\", [f.__name__ for f in app_methods]\n",
    "    )\n",
    "\n",
    "    self.mocks = self.AppMocks(  # type: ignore\n",
    "        **{\n",
    "            f.__name__: AsyncMock() if inspect.iscoroutinefunction(f) else MagicMock()\n",
    "            for f in app_methods\n",
    "        }\n",
    "    )\n",
    "\n",
    "    self.awaited_mocks = self.AppMocks(  # type: ignore\n",
    "        **{name: AwaitedMock(mock) for name, mock in self.mocks._asdict().items()}\n",
    "    )\n",
    "\n",
    "    def add_mock(\n",
    "        f: Callable[..., Any], mock: Union[AsyncMock, MagicMock]\n",
    "    ) -> Callable[..., Any]:\n",
    "        \"\"\"Add call to mock when calling function f\"\"\"\n",
    "\n",
    "        @functools.wraps(f)\n",
    "        async def async_inner(\n",
    "            *args: Any, f: Callable[..., Any] = f, mock: AsyncMock = mock, **kwargs: Any\n",
    "        ) -> Any:\n",
    "            await mock(*deepcopy(args), **kwargs)\n",
    "            return await f(*args, **kwargs)\n",
    "\n",
    "        @functools.wraps(f)\n",
    "        def sync_inner(\n",
    "            *args: Any, f: Callable[..., Any] = f, mock: MagicMock = mock, **kwargs: Any\n",
    "        ) -> Any:\n",
    "            mock(*deepcopy(args), **kwargs)\n",
    "            return f(*args, **kwargs)\n",
    "\n",
    "        if inspect.iscoroutinefunction(f):\n",
    "            return async_inner\n",
    "        else:\n",
    "            return sync_inner\n",
    "\n",
    "    self._consumers_store.update(\n",
    "        {\n",
    "            name: (\n",
    "                add_mock(f, getattr(self.mocks, f.__name__)),\n",
    "                decoder_fn,\n",
    "                executor,\n",
    "                kafka_brokers,\n",
    "                kwargs,\n",
    "            )\n",
    "            for name, (\n",
    "                f,\n",
    "                decoder_fn,\n",
    "                executor,\n",
    "                kafka_brokers,\n",
    "                kwargs,\n",
    "            ) in self._consumers_store.items()\n",
    "        }\n",
    "    )\n",
    "\n",
    "    self._producers_store.update(\n",
    "        {\n",
    "            name: (\n",
    "                add_mock(f, getattr(self.mocks, f.__name__)),\n",
    "                producer,\n",
    "                kafka_brokers,\n",
    "                kwargs,\n",
    "            )\n",
    "            for name, (\n",
    "                f,\n",
    "                producer,\n",
    "                kafka_brokers,\n",
    "                kwargs,\n",
    "            ) in self._producers_store.items()\n",
    "        }\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0ed851c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TestMsg(BaseModel):\n",
    "    msg: str = Field(...)\n",
    "\n",
    "\n",
    "app = FastKafka(kafka_brokers=dict(localhost=dict(url=\"localhost\", port=9092)))\n",
    "\n",
    "\n",
    "@app.consumes()\n",
    "async def on_preprocessed_signals(msg: TestMsg):\n",
    "    await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "@app.consumes(\n",
    "    topic=\"preprocessed_signals\",\n",
    "    brokers=dict(\n",
    "        localhost=[\n",
    "            dict(url=\"localhost\", port=9092),\n",
    "            dict(url=\"localhost\", port=9093),\n",
    "        ]\n",
    "    ),\n",
    ")\n",
    "async def on_preprocessed_signals_second(msg: TestMsg):\n",
    "    await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "@app.produces()\n",
    "async def to_predictions(prediction: TestMsg) -> TestMsg:\n",
    "    print(f\"Sending prediction: {prediction}\")\n",
    "    return prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca781c10",
   "metadata": {},
   "outputs": [],
   "source": [
    "app.create_mocks()\n",
    "app.mocks.on_preprocessed_signals.assert_not_awaited()\n",
    "app.mocks.on_preprocessed_signals_second.assert_not_awaited()\n",
    "app.mocks.to_predictions.assert_not_awaited()\n",
    "app.create_mocks()\n",
    "app.mocks.on_preprocessed_signals.assert_not_awaited()\n",
    "app.mocks.on_preprocessed_signals_second.assert_not_awaited()\n",
    "app.mocks.to_predictions.assert_not_awaited()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89b07708",
   "metadata": {},
   "outputs": [],
   "source": [
    "with pytest.raises(AssertionError) as e:\n",
    "    await app.awaited_mocks.on_preprocessed_signals.assert_called_with(123, timeout=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afa1957c",
   "metadata": {},
   "outputs": [],
   "source": [
    "app.create_mocks()\n",
    "app.mocks.on_preprocessed_signals.assert_not_awaited()\n",
    "await app.awaited_mocks.on_preprocessed_signals.assert_not_awaited(timeout=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bdaa582",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "def benchmark(\n",
    "    self: FastKafka,\n",
    "    interval: Union[int, timedelta] = 1,\n",
    "    *,\n",
    "    sliding_window_size: Optional[int] = None,\n",
    ") -> Callable[[Callable[[I], Optional[O]]], Callable[[I], Optional[O]]]:\n",
    "    \"\"\"Decorator to benchmark produces/consumes functions\n",
    "\n",
    "    Args:\n",
    "        interval: Period to use to calculate throughput. If value is of type int,\n",
    "            then it will be used as seconds. If value is of type timedelta,\n",
    "            then it will be used as it is. default: 1 - one second\n",
    "        sliding_window_size: The size of the sliding window to use to calculate\n",
    "            average throughput. default: None - By default average throughput is\n",
    "            not calculated\n",
    "    \"\"\"\n",
    "\n",
    "    def _decorator(func: Callable[[I], Optional[O]]) -> Callable[[I], Optional[O]]:\n",
    "        func_name = f\"{func.__module__}.{func.__qualname__}\"\n",
    "\n",
    "        @wraps(func)\n",
    "        def wrapper(\n",
    "            *args: I,\n",
    "            **kwargs: I,\n",
    "        ) -> Optional[O]:\n",
    "            _benchmark(\n",
    "                interval=interval,\n",
    "                sliding_window_size=sliding_window_size,\n",
    "                func_name=func_name,\n",
    "                benchmark_results=self.benchmark_results,\n",
    "            )\n",
    "            return func(*args, **kwargs)\n",
    "\n",
    "        @wraps(func)\n",
    "        async def async_wrapper(\n",
    "            *args: I,\n",
    "            **kwargs: I,\n",
    "        ) -> Optional[O]:\n",
    "            _benchmark(\n",
    "                interval=interval,\n",
    "                sliding_window_size=sliding_window_size,\n",
    "                func_name=func_name,\n",
    "                benchmark_results=self.benchmark_results,\n",
    "            )\n",
    "            return await func(*args, **kwargs)  # type: ignore\n",
    "\n",
    "        if inspect.iscoroutinefunction(func):\n",
    "            return async_wrapper  # type: ignore\n",
    "        else:\n",
    "            return wrapper\n",
    "\n",
    "    return _decorator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c4fcda2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "Hello I am over after 10k msgs\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "Hello I am over after 10k msgs\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "for executor in [\"SequentialExecutor\", \"DynamicTaskExecutor\"]:\n",
    "\n",
    "    class TestMsg(BaseModel):\n",
    "        msg: str = Field(...)\n",
    "\n",
    "    app = FastKafka(kafka_brokers=dict(localhost=dict(url=\"localhost\", port=9092)))\n",
    "    # app.benchmark_results[\"test\"] = dict(count=0)\n",
    "\n",
    "    @app.consumes(executor=executor)\n",
    "    #@app.benchmark(interval=1, sliding_window_size=5)\n",
    "    async def on_preprocessed_signals(msg: TestMsg):\n",
    "        await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "    @app.produces()\n",
    "    #@app.benchmark(interval=1, sliding_window_size=5)\n",
    "    async def to_predictions(prediction: TestMsg) -> TestMsg:\n",
    "        #         print(f\"Sending prediction: {prediction}\")\n",
    "        return prediction\n",
    "\n",
    "    async with Tester(app) as tester:\n",
    "        for i in range(10_000):\n",
    "            await tester.to_preprocessed_signals(TestMsg(msg=f\"signal {i}\"))\n",
    "        print(\"Hello I am over after 10k msgs\")\n",
    "        await asyncio.sleep(5)\n",
    "        tester.mocks.on_predictions.assert_called()\n",
    "\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d213e250",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "def fastapi_lifespan(\n",
    "    self: FastKafka, kafka_broker_name: str\n",
    ") -> Callable[[\"FastAPI\"], AsyncIterator[None]]:\n",
    "    \"\"\"\n",
    "    Method for managing the lifespan of a FastAPI application with a specific Kafka broker.\n",
    "\n",
    "    Args:\n",
    "        kafka_broker_name: The name of the Kafka broker to start FastKafka\n",
    "\n",
    "    Returns:\n",
    "        Lifespan function to use for initializing FastAPI\n",
    "    \"\"\"\n",
    "\n",
    "    @asynccontextmanager\n",
    "    async def lifespan(fastapi_app: \"FastAPI\") -> AsyncIterator[None]:\n",
    "        self.set_kafka_broker(kafka_broker_name=kafka_broker_name)\n",
    "        async with self:\n",
    "            yield\n",
    "\n",
    "    return lifespan  # type: ignore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b51932c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TestMsg(BaseModel):\n",
    "    msg: str = Field(...)\n",
    "\n",
    "\n",
    "app_for_tester = FastKafka(\n",
    "    kafka_brokers=dict(localhost=dict(url=\"localhost\", port=9092)),\n",
    "    group_id=\"app_for_tester_group\",\n",
    ")\n",
    "\n",
    "\n",
    "@app_for_tester.consumes(topic=\"preprocessed_signals\")\n",
    "async def on_app_for_tester_preprocessed_signals(msg: TestMsg):\n",
    "    print(\"receving messages on app_for_tester\")\n",
    "    await to_app_for_tester_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "@app_for_tester.produces(topic=\"predictions\")\n",
    "async def to_app_for_tester_predictions(prediction: TestMsg) -> TestMsg:\n",
    "    print(\"sending predictions on app_for_tester\")\n",
    "    return prediction\n",
    "\n",
    "\n",
    "def create_app_for_fastapi(port: int):\n",
    "    app_for_fastapi = FastKafka(\n",
    "        kafka_brokers=dict(localhost=dict(url=\"localhost\", port=port)),\n",
    "        group_id=\"app_for_fastapi_group\",\n",
    "    )\n",
    "\n",
    "    @app_for_fastapi.consumes(topic=\"preprocessed_signals\")\n",
    "    async def on_app_for_fastapi_preprocessed_signals(msg: TestMsg):\n",
    "        print(\"receving messages on app_for_fastapi\")\n",
    "        await to_app_for_fastapi_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "    @app_for_fastapi.produces(topic=\"predictions\")\n",
    "    async def to_app_for_fastapi_predictions(prediction: TestMsg) -> TestMsg:\n",
    "        print(\"sending predictions on app_for_fastapi\")\n",
    "        return prediction\n",
    "\n",
    "    return app_for_fastapi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e753c66",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'group_id': 'app_for_tester_group', 'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:     Started server process [8091]\n",
      "INFO:     Waiting for application startup.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'group_id': 'app_for_fastapi_group', 'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:     Application startup complete.\n",
      "INFO:     Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:     127.0.0.1:35572 - \"GET /predict HTTP/1.1\" 200 OK\n",
      "receving messages on app_for_tester\n",
      "sending predictions on app_for_tester\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:     Shutting down\n",
      "INFO:     Waiting for application shutdown.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:     Application shutdown complete.\n",
      "INFO:     Finished server process [8091]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exception raised e=AssertionError(1)\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "def run_uvicorn():\n",
    "    fastapi_app = FastAPI(lifespan=app_for_fastapi.fastapi_lifespan(\"localhost\"))\n",
    "\n",
    "    @fastapi_app.get(\"/predict\")\n",
    "    async def predict():\n",
    "        return {\"result\": \"hello\"}\n",
    "\n",
    "    uvicorn.run(\n",
    "        fastapi_app,\n",
    "        host=\"0.0.0.0\",\n",
    "        port=8000,\n",
    "        reload=False,\n",
    "        log_level=\"debug\",\n",
    "        workers=1,\n",
    "    )\n",
    "\n",
    "\n",
    "async with Tester(app_for_tester) as tester:\n",
    "    app_for_fastapi = create_app_for_fastapi(port=9092)\n",
    "    with run_in_process(run_uvicorn) as p:\n",
    "        await asyncio.sleep(3)\n",
    "        res = requests.get(\"http://127.0.0.1:8000/predict\")\n",
    "        assert res.ok\n",
    "\n",
    "        await tester.to_preprocessed_signals(TestMsg(msg=f\"signal 10\"))\n",
    "        await asyncio.sleep(3)\n",
    "        assert (\n",
    "            tester.mocks.on_predictions.call_count == 2\n",
    "        ), tester.mocks.on_predictions.call_count\n",
    "\n",
    "    p.close()\n",
    "\n",
    "print(\"ok\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
